00001
00002
00003 import string
00004 from random import random, shuffle
00005
00006 import os
00007
00008 import pygraphviz as pgv
00009
00010 from time import sleep
00011
00012 class TaskStatus(object):
00013 """ A class for enumerating task statuses """
00014 FAILURE = 0
00015 SUCCESS = 1
00016 RUNNING = 2
00017
00018
00019 last_dot_tree = ''
00020
00021 def weighted_choice(weights):
00022 rnd = random() * sum(weights)
00023 for i, w in enumerate(weights):
00024 rnd -= w
00025 if rnd < 0:
00026 return i
00027
00028 def weighted_shuffle(a,w):
00029 w = list(w)
00030 if len(a) != len(w):
00031 print("ERROR: Weighted_shuffle: Number of weights does not match number of children.")
00032 os._exit(0)
00033 return
00034
00035 r = [0]*len(a)
00036 for i in range(len(a)):
00037 j = weighted_choice(w)
00038 r[i]=a[j]
00039 w[j] = 0
00040 return r
00041
00042 class Task(object):
00043 """ "The base Task class """
00044 def __init__(self, name, children=None, reset_after=False, announce=False, *args, **kwargs):
00045 self.name = name
00046 self.status = None
00047 self.reset_after = reset_after
00048 self._announce = announce
00049
00050 if children is None:
00051 children = []
00052
00053 self.children = children
00054
00055 def run(self):
00056 pass
00057
00058 def reset(self):
00059 for c in self.children:
00060 c.reset()
00061
00062 self.status = None
00063
00064 def add_child(self, c):
00065 self.children.append(c)
00066
00067 def remove_child(self, c):
00068 self.children.remove(c)
00069
00070 def prepend_child(self, c):
00071 self.children.insert(0, c)
00072
00073 def insert_child(self, c, i):
00074 self.children.insert(i, c)
00075
00076 def get_status(self):
00077 return self.status
00078
00079 def set_status(self, s):
00080 self.status = s
00081
00082 def announce(self):
00083 print("Executing task " + str(self.name))
00084
00085 def get_type(self):
00086 return type(self)
00087
00088
00089 def __enter__(self):
00090 return self.name
00091
00092 def __exit__(self, exc_type, exc_val, exc_tb):
00093 if exc_type is not None:
00094 return False
00095 return True
00096
00097 class Selector(Task):
00098 """ A selector runs each task in order until one succeeds,
00099 at which point it returns SUCCESS. If all tasks fail, a FAILURE
00100 status is returned. If a subtask is still RUNNING, then a RUNNING
00101 status is returned and processing continues until either SUCCESS
00102 or FAILURE is returned from the subtask.
00103 """
00104 def __init__(self, name, *args, **kwargs):
00105 super(Selector, self).__init__(name, *args, **kwargs)
00106
00107 def run(self):
00108 for c in self.children:
00109
00110 c.status = c.run()
00111
00112 if c.status != TaskStatus.FAILURE:
00113 if c.status == TaskStatus.SUCCESS:
00114 if self.reset_after:
00115 self.reset()
00116 return TaskStatus.SUCCESS
00117 else:
00118 return c.status
00119 return c.status
00120
00121 if self.reset_after:
00122 self.reset()
00123
00124 return TaskStatus.FAILURE
00125
00126 class Sequence(Task):
00127 """
00128 A sequence runs each task in order until one fails,
00129 at which point it returns FAILURE. If all tasks succeed, a SUCCESS
00130 status is returned. If a subtask is still RUNNING, then a RUNNING
00131 status is returned and processing continues until either SUCCESS
00132 or FAILURE is returned from the subtask.
00133 """
00134 def __init__(self, name, *args, **kwargs):
00135 super(Sequence, self).__init__(name, *args, **kwargs)
00136
00137 def run(self):
00138 if self._announce:
00139 self.announce()
00140
00141 for c in self.children:
00142
00143 c.status = c.run()
00144
00145 if c.status != TaskStatus.SUCCESS:
00146 if c.status == TaskStatus.FAILURE:
00147 if self.reset_after:
00148 self.reset()
00149 return TaskStatus.FAILURE
00150 return c.status
00151
00152 if self.reset_after:
00153 self.reset()
00154
00155 return TaskStatus.SUCCESS
00156
00157 class RandomSelector(Task):
00158 """ A selector runs each task in random order until one succeeds,
00159 at which point it returns SUCCESS. If all tasks fail, a FAILURE
00160 status is returned. If a subtask is still RUNNING, then a RUNNING
00161 status is returned and processing continues until either SUCCESS
00162 or FAILURE is returned from the subtask.
00163 """
00164 def __init__(self, name, *args, **kwargs):
00165 super(RandomSelector, self).__init__(name, *args, **kwargs)
00166
00167 self.shuffled = False
00168
00169 def run(self):
00170 if not self.shuffled:
00171 shuffle(self.children)
00172 self.shuffled = True
00173
00174 for c in self.children:
00175
00176 c.status = c.run()
00177
00178 if c.status != TaskStatus.FAILURE:
00179 if c.status == TaskStatus.SUCCESS:
00180 self.shuffled = False
00181
00182 return c.status
00183
00184 self.shuffled = False
00185
00186 if self.reset_after:
00187 self.reset()
00188
00189 return TaskStatus.FAILURE
00190
00191 class RandomSequence(Task):
00192 """
00193 A sequence runs each task in random order until one fails,
00194 at which point it returns FAILURE. If all tasks succeed, a SUCCESS
00195 status is returned. If a subtask is still RUNNING, then a RUNNING
00196 status is returned and processing continues until either SUCCESS
00197 or FAILURE is returned from the subtask.
00198 """
00199 def __init__(self, name, *args, **kwargs):
00200 super(RandomSequence, self).__init__(name, *args, **kwargs)
00201
00202 self.shuffled = False
00203
00204 def run(self):
00205 if not self.shuffled:
00206 shuffle(self.children)
00207 self.shuffled = True
00208
00209 for c in self.children:
00210
00211 c.status = c.run()
00212
00213 if c.status != TaskStatus.SUCCESS:
00214 if c.status == TaskStatus.FAILURE:
00215 self.shuffled = False
00216
00217 return c.status
00218
00219 self.shuffled = False
00220
00221 if self.reset_after:
00222 self.reset()
00223
00224 return TaskStatus.SUCCESS
00225
00226 class WeightedRandomSequence(Task):
00227 """
00228 A sequence runs each task in random order until one fails,
00229 at which point it returns FAILURE. If all tasks succeed, a SUCCESS
00230 status is returned. If a subtask is still RUNNING, then a RUNNING
00231 status is returned and processing continues until either SUCCESS
00232 or FAILURE is returned from the subtask.
00233 """
00234 def __init__(self, name, weights, *args, **kwargs):
00235 super(WeightedRandomSequence, self).__init__(name, *args, **kwargs)
00236
00237 self.shuffled = False
00238 self.weights = weights
00239
00240 def run(self):
00241 if not self.shuffled:
00242 shuffled_children = weighted_shuffle(self.children, self.weights)
00243 self.shuffled = True
00244
00245 for c in shuffled_children:
00246
00247 c.status = c.run()
00248
00249 if c.status != TaskStatus.SUCCESS:
00250 if c.status == TaskStatus.FAILURE:
00251 self.shuffled = False
00252
00253 return c.status
00254
00255 self.shuffled = False
00256
00257 if self.reset_after:
00258 self.reset()
00259
00260 return TaskStatus.SUCCESS
00261
00262 class WeightedRandomSelector(Task):
00263 """ A selector runs each task in random order until one succeeds,
00264 at which point it returns SUCCESS. If all tasks fail, a FAILURE
00265 status is returned. If a subtask is still RUNNING, then a RUNNING
00266 status is returned and processing continues until either SUCCESS
00267 or FAILURE is returned from the subtask.
00268 """
00269 def __init__(self, name, weights, *args, **kwargs):
00270 super(WeightedRandomSelector, self).__init__(name, *args, **kwargs)
00271
00272 self.shuffled = False
00273 self.weights = weights
00274
00275 def run(self):
00276 if not self.shuffled:
00277 shuffled_children = weighted_shuffle(self.children, self.weights)
00278 self.shuffled = True
00279
00280 for c in shuffled_children:
00281
00282 c.status = c.run()
00283
00284 if c.status != TaskStatus.FAILURE:
00285 if c.status == TaskStatus.SUCCESS:
00286 self.shuffled = False
00287
00288 return c.status
00289
00290 self.shuffled = False
00291
00292 if self.reset_after:
00293 self.reset()
00294
00295 return TaskStatus.FAILURE
00296
00297 class Iterator(Task):
00298 """
00299 Iterate through all child tasks ignoring failure.
00300 """
00301 def __init__(self, name, *args, **kwargs):
00302 super(Iterator, self).__init__(name, *args, **kwargs)
00303
00304 def run(self):
00305 for c in self.children:
00306
00307 c.status = c.run()
00308
00309 if c.status == TaskStatus.RUNNING:
00310 return c.status
00311
00312 if self.reset_after:
00313 self.reset()
00314
00315 return TaskStatus.SUCCESS
00316
00317 class RandomIterator(Task):
00318 """
00319 Iterate through all child tasks randomly (without replacement) ignoring failure.
00320 """
00321 def __init__(self, name, *args, **kwargs):
00322 super(RandomIterator, self).__init__(name, *args, **kwargs)
00323
00324 self.shuffled = False
00325
00326 def run(self):
00327 if not self.shuffled:
00328 shuffle(self.children)
00329 self.shuffled = True
00330
00331 for c in self.children:
00332
00333 c.status = c.run()
00334
00335 if c.status == TaskStatus.RUNNING:
00336 return c.status
00337
00338 self.shuffled = False
00339
00340 if self.reset_after:
00341 self.reset()
00342
00343 return TaskStatus.SUCCESS
00344
00345 class WeightedRandomIterator(Task):
00346 """
00347 Iterate through all child tasks randomly (without replacement) ignoring failure.
00348 """
00349 def __init__(self, name, weights, *args, **kwargs):
00350 super(WeightedRandomIterator, self).__init__(name, *args, **kwargs)
00351
00352 self.shuffled = False
00353 self.weights = weights
00354
00355 def run(self):
00356 if not self.shuffled:
00357 suffled_children = weighted_shuffle(self.children, self.weights)
00358 self.shuffled = True
00359
00360 for c in suffled_children:
00361
00362 c.status = c.run()
00363
00364 if c.status == TaskStatus.RUNNING:
00365 return c.status
00366
00367 self.shuffled = False
00368
00369 if self.reset_after:
00370 self.reset()
00371
00372 return TaskStatus.SUCCESS
00373
00374 class ParallelOne(Task):
00375 """
00376 A parallel task runs each child task at (roughly) the same time.
00377 The ParallelOne task returns success as soon as any child succeeds.
00378 """
00379 def __init__(self, name, *args, **kwargs):
00380 super(ParallelOne, self).__init__(name, *args, **kwargs)
00381
00382 self.failure = dict()
00383 self.index = 0
00384
00385 def run(self):
00386 n_children = len(self.children)
00387
00388 if self.index < n_children:
00389 child = self.children[self.index]
00390 child.status = child.run()
00391
00392 if child.status != TaskStatus.SUCCESS:
00393 self.index += 1
00394
00395 if child.status == TaskStatus.FAILURE:
00396 self.failure[child.name] = TaskStatus.FAILURE
00397
00398 return TaskStatus.RUNNING
00399 else:
00400 if self.reset_after:
00401 self.reset()
00402 self.index = 0
00403 return TaskStatus.SUCCESS
00404
00405 elif len(self.failure) == n_children:
00406 if self.reset_after:
00407 self.reset()
00408 return TaskStatus.FAILURE
00409 else:
00410 self.index = 0
00411 return TaskStatus.RUNNING
00412
00413 def reset(self):
00414 super(ParallelOne, self).reset()
00415 self.failure = dict()
00416 self.index = 0
00417
00418 class ParallelAll(Task):
00419 """
00420 A parallel task runs each child task at (roughly) the same time.
00421 The ParallelAll task requires all subtasks to succeed for it to succeed.
00422 """
00423 def __init__(self, name, *args, **kwargs):
00424 super(ParallelAll, self).__init__(name, *args, **kwargs)
00425
00426 self.success = dict()
00427 self.index = 0
00428
00429 def run(self):
00430 n_children = len(self.children)
00431
00432 if self.index < n_children:
00433 child = self.children[self.index]
00434 child.status = child.run()
00435
00436 if child.status != TaskStatus.FAILURE:
00437 self.index += 1
00438
00439 if child.status == TaskStatus.SUCCESS:
00440 self.success[child.name] = TaskStatus.SUCCESS
00441
00442 return TaskStatus.RUNNING
00443 else:
00444 if self.reset_after:
00445 self.reset()
00446 return TaskStatus.FAILURE
00447
00448 elif len(self.success) == n_children:
00449 if self.reset_after:
00450 self.reset()
00451 return TaskStatus.SUCCESS
00452 else:
00453 self.index = 0
00454 return TaskStatus.RUNNING
00455
00456 def reset(self):
00457 super(ParallelAll, self).reset()
00458 self.success = dict()
00459 self.index = 0
00460
00461 class Loop(Task):
00462 """
00463 Loop over one or more subtasks for the given number of iterations
00464 Use the value -1 to indicate a continual loop.
00465 """
00466 def __init__(self, name, announce=True, *args, **kwargs):
00467 super(Loop, self).__init__(name, *args, **kwargs)
00468
00469 self.iterations = kwargs['iterations']
00470 self._announce = announce
00471 self.loop_count = 0
00472 self.name = name
00473 print("Loop iterations: " + str(self.iterations))
00474
00475 def run(self):
00476
00477 c = self.children[0]
00478
00479 while self.iterations == -1 or self.loop_count < self.iterations:
00480 c.status = c.run()
00481
00482 status = c.status
00483
00484 self.status = status
00485
00486 if status == TaskStatus.SUCCESS or status == TaskStatus.FAILURE:
00487 self.loop_count += 1
00488
00489 if self._announce:
00490 print(self.name + " COMPLETED " + str(self.loop_count) + " LOOP(S)")
00491
00492 c.reset()
00493
00494 return status
00495
00496 class Limit(Task):
00497 """
00498 Limit the number of times a task can execute
00499 """
00500 def __init__(self, name, announce=True, *args, **kwargs):
00501 super(Limit, self).__init__(name, *args, **kwargs)
00502
00503 self.max_executions = kwargs['max_executions']
00504 self._announce = announce
00505 self.execution_count = 0
00506 self.name = name
00507 print("Limit number of executions to: " + str(self.max_executions))
00508
00509 def run(self):
00510 c = self.children[0]
00511
00512 if self.execution_count >= self.max_executions:
00513
00514 if self._announce:
00515 print(self.name + " reached maximum number (" + str(self.max_executions) + ") of executions.")
00516
00517 if self.reset_after:
00518 self.reset()
00519
00520 return TaskStatus.FAILURE
00521
00522 c.status = c.run()
00523
00524 if c.status != TaskStatus.RUNNING:
00525 self.execution_count += 1
00526
00527 return c.status
00528
00529 def reset(self):
00530 c = self.children[0]
00531 c.reset()
00532 self.execution_count = 0
00533 self.status = None
00534
00535 class IgnoreFailure(Task):
00536 """
00537 Always return either RUNNING or SUCCESS.
00538 """
00539 def __init__(self, name, *args, **kwargs):
00540 super(IgnoreFailure, self).__init__(name, *args, **kwargs)
00541
00542 def run(self):
00543
00544 for c in self.children:
00545
00546 c.status = c.run()
00547
00548 if c.status == TaskStatus.FAILURE:
00549 return TaskStatus.SUCCESS
00550 else:
00551 return c.status
00552
00553 return TaskStatus.SUCCESS
00554
00555
00556 class AlwaysFail(Task):
00557 """
00558 Always return FAILURE
00559 """
00560 def __init__(self, name, *args, **kwargs):
00561 super(AlwaysFail, self).__init__(name, *args, **kwargs)
00562
00563 def run(self):
00564
00565 for c in self.children:
00566
00567 c.status = c.run()
00568
00569 return c.status
00570
00571 return TaskStatus.FAILURE
00572
00573 class AlwaysSucceed(Task):
00574 """
00575 Always return SUCCESS
00576 """
00577 def __init__(self, name, *args, **kwargs):
00578 super(AlwaysSucceed, self).__init__(name, *args, **kwargs)
00579
00580 def run(self):
00581
00582 for c in self.children:
00583
00584 c.status = c.run()
00585
00586 return c.status
00587
00588 return TaskStatus.SUCCESS
00589
00590
00591 class Invert(Task):
00592 """
00593 Turn SUCCESS into FAILURE and vice-versa
00594 """
00595 def __init__(self, name, *args, **kwargs):
00596 super(Invert, self).__init__(name, *args, **kwargs)
00597
00598 def run(self):
00599
00600 for c in self.children:
00601
00602 c.status = c.run()
00603
00604 if c.status == TaskStatus.FAILURE:
00605 return TaskStatus.SUCCESS
00606
00607 elif c.status == TaskStatus.SUCCESS:
00608 return TaskStatus.FAILURE
00609
00610 else:
00611 return c.status
00612
00613
00614 TaskNot = Invert
00615
00616 class UntilFailure(Task):
00617 """
00618 Continue executing a task until it fails or max_attempts is reached.
00619 """
00620 def __init__(self, name, max_attempts=-1, *args, **kwargs):
00621 super(UntilFailure, self).__init__(name, *args, **kwargs)
00622
00623 self.max_attempts = max_attempts
00624 self.attempts = 0
00625
00626 def run(self):
00627 c = self.children[0]
00628
00629 while self.attempts < self.max_attempts or self.max_attempts == -1:
00630
00631 c.status = c.run()
00632
00633 if c.status == TaskStatus.FAILURE:
00634 return TaskStatus.SUCCESS
00635
00636 if c.status == TaskStatus.SUCCESS:
00637 self.attempts += 1
00638 c.reset()
00639
00640 return TaskStatus.RUNNING
00641
00642 return TaskStatus.FAILURE
00643
00644 def reset(self):
00645 c = self.children[0]
00646 c.reset()
00647 self.attempts = 0
00648
00649
00650 UntilFail = UntilFailure
00651
00652 class UntilSuccess(Task):
00653 """
00654 Continue executing a task until it succeeds or max_attempts is reached.
00655 """
00656 def __init__(self, name, max_attempts=-1, *args, **kwargs):
00657 super(UntilSuccess, self).__init__(name, *args, **kwargs)
00658
00659 self.max_attempts = max_attempts
00660 self.attempts = 0
00661
00662 def run(self):
00663 c = self.children[0]
00664
00665 while self.attempts < self.max_attempts or self.max_attempts == -1:
00666
00667 c.status = c.run()
00668
00669 if c.status == TaskStatus.SUCCESS:
00670 return TaskStatus.SUCCESS
00671
00672 if c.status == TaskStatus.FAILURE:
00673 self.attempts += 1
00674 c.reset()
00675
00676 return TaskStatus.RUNNING
00677
00678 if self.reset_after:
00679 self.reset()
00680
00681 return TaskStatus.FAILURE
00682
00683 def reset(self):
00684 c = self.children[0]
00685 c.reset()
00686 self.attempts = 0
00687
00688 class AutoRemoveSequence(Task):
00689 """
00690 Remove each successful subtask from a sequence
00691 """
00692 def __init__(self, name, *args, **kwargs):
00693 super(AutoRemoveSequence, self).__init__(name, *args, **kwargs)
00694
00695 def run(self):
00696 for c in self.children:
00697 c.status = c.run()
00698
00699 if c.status == TaskStatus.FAILURE:
00700 return TaskStatus.FAILURE
00701
00702 if c.statuss == TaskStatus.RUNNING:
00703 return TaskStatus.RUNNING
00704
00705 try:
00706 self.children.remove(self.children[0])
00707 except:
00708 return TaskStatus.FAILURE
00709
00710 return TaskStatus.SUCCESS
00711
00712 class CallbackTask(Task):
00713 """
00714 Turn any callback function (cb) into a task
00715 """
00716 def __init__(self, name, cb=None, cb_args=[], cb_kwargs={}, **kwargs):
00717 super(CallbackTask, self).__init__(name, cb=None, cb_args=[], cb_kwargs={}, **kwargs)
00718
00719 self.name = name
00720 self.cb = cb
00721 self.cb_args = cb_args
00722 self.cb_kwargs = cb_kwargs
00723
00724 def run(self):
00725 status = self.cb(*self.cb_args, **self.cb_kwargs)
00726
00727 if status is None:
00728 self.status = TaskStatus.RUNNING
00729
00730 elif status:
00731 self.status = TaskStatus.SUCCESS
00732
00733 else:
00734 self.status = TaskStatus.FAILURE
00735
00736 return self.status
00737
00738 def reset(self):
00739 self.status = None
00740
00741 class WaitTask(Task):
00742 """
00743 This is a *blocking* wait task. The interval argument is in seconds.
00744 """
00745 def __init__(self, name, interval, *args, **kwargs):
00746 super(WaitTask, self).__init__(name, interval, *args, **kwargs)
00747
00748 self.interval = interval
00749
00750 def run(self):
00751 sleep(self.interval)
00752
00753 return TaskStatus.SUCCESS
00754
00755 class loop(Task):
00756 """
00757 Loop over one or more subtasks a given number of iterations
00758 """
00759 def __init__(self, task, iterations=-1):
00760 new_name = task.name + "_loop_" + str(iterations)
00761 super(loop, self).__init__(new_name)
00762
00763 self.iterations = iterations
00764 self.old_run = task.run
00765 self.old_reset = task.reset
00766 self.old_children = task.children
00767 self.loop_count = 0
00768
00769 print("Loop iterations: " + str(self.iterations))
00770
00771 def run(self):
00772 if self.iterations != -1 and self.loop_count >= self.iterations:
00773 return TaskStatus.SUCCESS
00774
00775 print("Loop " + str(self.loop_count + 1))
00776
00777 while True:
00778 self.status = self.old_run()
00779
00780 if self.status == TaskStatus.SUCCESS:
00781 break
00782 else:
00783 return self.status
00784
00785 self.old_reset()
00786 self.loop_count += 1
00787
00788 class limit(Task):
00789 """
00790 Limit a task to the given number of executions
00791 """
00792 def __init__(self, task, max_executions=-1):
00793 new_name = task.name + "_limit_" + str(max_executions)
00794 super(limit, self).__init__(new_name)
00795
00796 self.max_executions = max_executions
00797 self.old_run = task.run
00798 self.old_reset = task.reset
00799 self.old_children = task.children
00800 self.execution_count = 0
00801
00802 print("Limit number of executions to: " + str(self.max_executions))
00803
00804 def run(self):
00805 if self.max_executions != -1 and self.execution_count >= self.max_executions:
00806 self.execution_count = 0
00807
00808 if self._announce:
00809 print(self.name + " reached maximum number (" + str(self.max_executions) + ") of executions.")
00810
00811 return TaskStatus.FAILURE
00812
00813 while True:
00814 self.status = self.old_run()
00815
00816 if self.status == TaskStatus.SUCCESS:
00817 break
00818 else:
00819 return self.status
00820
00821 self.old_reset()
00822 self.execution_count += 1
00823
00824 class ignore_failure(Task):
00825 """
00826 Always return either RUNNING or SUCCESS.
00827 """
00828 def __init__(self, task):
00829 new_name = task.name + "_ignore_failure"
00830 super(ignore_failure, self).__init__(new_name)
00831
00832 self.old_run = task.run
00833
00834 def run(self):
00835 while True:
00836 self.status = self.old_run()
00837
00838 if self.status == TaskStatus.FAILURE:
00839 return TaskStatus.SUCCESS
00840 else:
00841 return self.status
00842
00843 class ignore_success(Task):
00844 """
00845 Always return FAILURE or RUNNING
00846 """
00847 def __init__(self, task):
00848 new_name = task.name + "_ignore_success"
00849 super(ignore_success, self).__init__(new_name)
00850
00851 self.old_run = task.run
00852
00853 def run(self):
00854 while True:
00855 self.status = self.old_run()
00856
00857 if self.status == TaskStatus.SUCCESS:
00858 return TaskStatus.FAILURE
00859 else:
00860 return self.status
00861
00862 class task_not(Task):
00863 """
00864 Turn SUCCESS into FAILURE and vice-versa
00865 """
00866 def __init__(self, task):
00867 new_name = task.name + "_not"
00868 super(task_not, self).__init__(new_name)
00869
00870 self.old_run = task.run
00871
00872 def run(self):
00873 while True:
00874 self.status = self.old_run()
00875
00876 if self.status == TaskStatus.FAILURE:
00877 return TaskStatus.SUCCESS
00878
00879 elif self.status == TaskStatus.SUCCESS:
00880 return TaskStatus.FAILURE
00881
00882 else:
00883 return self.status
00884
00885
00886 invert = task_not
00887
00888 class until_fail(Task):
00889 """
00890 Execute a task until it fails
00891 """
00892 def __init__(self, task):
00893 new_name = task.name + "_until_fail"
00894 super(until_fail, self).__init__(new_name)
00895
00896 self.old_run = task.run
00897
00898 def run(self):
00899 while True:
00900 self.status = self.old_run()
00901
00902 if self.status == TaskStatus.FAILURE:
00903 break
00904
00905 else:
00906 return self.status
00907
00908 return TaskStatus.SUCCESS
00909
00910 class always_fail(Task):
00911 """
00912 Execute a task but always return FAILTURE
00913 """
00914 def __init__(self, task):
00915 new_name = task.name + "_always_fail"
00916 super(always_fail, self).__init__(new_name)
00917
00918 self.old_run = task.run
00919
00920 def run(self):
00921 while True:
00922 self.old_run()
00923
00924 self.status = TaskStatus.FAILURE
00925
00926 return TaskStatus.FAILURE
00927
00928
00929 def print_tree(tree, indent=0, use_symbols=False):
00930 """
00931 Print an ASCII representation of the tree
00932 """
00933 if use_symbols:
00934 if indent == 0:
00935 print_tree_symbol(tree, indent)
00936 indent += 1
00937
00938 for c in tree.children:
00939 print_tree_symbol(c, indent)
00940
00941 try:
00942 if c.children != []:
00943 print_tree(c, indent+1, use_symbols)
00944 except:
00945 pass
00946 else:
00947 for c in tree.children:
00948 print " " * indent, "-->", c.name
00949
00950 try:
00951 if c.children != []:
00952 print_tree(c, indent + 1)
00953 except:
00954 pass
00955
00956 def print_tree_symbol(c, indent):
00957 """
00958 Use ASCII symbols to represent Sequence, Selector, Task, etc.
00959 """
00960 if isinstance(c, Selector):
00961 print " " * indent, "--?",
00962 elif isinstance(c, (Sequence, Iterator)):
00963 print " " * indent, "-->",
00964 elif isinstance(c, (RandomSequence, RandomIterator, WeightedRandomSequence, WeightedRandomIterator)):
00965 print " " * indent, "~~>",
00966 elif isinstance(c, (RandomSelector, WeightedRandomSelector)):
00967 print " " * indent, "~~?",
00968 elif isinstance(c, ParallelOne):
00969 print " " * indent, "==?",
00970 elif isinstance(c, ParallelAll):
00971 print " " * indent, "==>",
00972 elif isinstance(c, Loop):
00973 print " " * indent, "<->",
00974 elif isinstance(c, Invert):
00975 print " " * indent, "--!",
00976 else:
00977 print " " * indent, "--|",
00978
00979 print c.name
00980
00981 def print_phpsyntax_tree(tree):
00982 """
00983 Print an output compatible with ironcreek.net/phpSyntaxTree
00984 """
00985 for c in tree.children:
00986 print "[" + string.replace(c.name, "_", "."),
00987 if c.children != []:
00988 print_phpsyntax_tree(c),
00989 print "]",
00990
00991 def print_dot_tree(root, dotfilepath=None):
00992 """
00993 Print an output compatible with the DOT synatax and Graphiz
00994 """
00995 gr = pgv.AGraph(strict=True, directed=True, rotate='0', bgcolor='white', ordering="out")
00996 gr.node_attr['fontsize'] = '9'
00997 gr.node_attr['color'] = 'black'
00998
00999 if dotfilepath is None:
01000 dotfilepath = os.path.expanduser('~') + '/.ros/tree.dot'
01001
01002 global last_dot_tree
01003
01004
01005 gr.add_node(root.name)
01006 node = gr.get_node(root.name)
01007
01008 if root.status == TaskStatus.RUNNING:
01009 node.attr['fillcolor'] = 'yellow'
01010 node.attr['style'] = 'filled'
01011 node.attr['border'] = 'bold'
01012 elif root.status == TaskStatus.SUCCESS:
01013 node.attr['color'] = 'green'
01014 elif root.status == TaskStatus.FAILURE:
01015 node.attr['color'] = 'red'
01016 else:
01017 node.attr['color'] = 'black'
01018
01019 def add_edges(root):
01020 for c in root.children:
01021 if isinstance(c, (Sequence, Iterator, RandomSequence, RandomIterator, WeightedRandomSequence, WeightedRandomIterator)):
01022 gr.add_node(c.name, shape="cds")
01023 elif isinstance(c, (Selector, RandomSelector, WeightedRandomSelector)):
01024 gr.add_node(c.name, shape="diamond")
01025 elif isinstance(c, (ParallelOne, ParallelAll)):
01026 gr.add_node(c.name, shape="parallelogram")
01027 elif isinstance(c, Invert):
01028 gr.add_node(c.name, shape="house")
01029 else:
01030 gr.add_node(c.name)
01031
01032 gr.add_edge((root.name, c.name))
01033 node = gr.get_node(c.name)
01034
01035 if c.status == TaskStatus.RUNNING:
01036 node.attr['fillcolor'] = 'yellow'
01037 node.attr['style'] = 'filled'
01038 node.attr['border'] = 'bold'
01039 elif c.status == TaskStatus.SUCCESS:
01040 node.attr['color'] = 'green'
01041 elif c.status == TaskStatus.FAILURE:
01042 node.attr['color'] = 'red'
01043 else:
01044 node.attr['color'] = 'black'
01045
01046 if c.children != []:
01047 add_edges(c)
01048
01049 add_edges(root)
01050
01051 current_dot_tree = gr.string()
01052
01053 if current_dot_tree != last_dot_tree:
01054 gr.write(dotfilepath)
01055 last_dot_tree = gr.string()