Tweaks code to auto tag

DAGWorks-Inc · Mar 7, 2024 · 4800e6d · 4800e6d
1 parent f8bf58d
commit 4800e6d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/examples/cmdline_orchestrator/cmdline.py b/examples/cmdline_orchestrator/cmdline.py
@@ -3,6 +3,7 @@
 
 from hamilton.execution.executors import DefaultExecutionManager, TaskExecutor
 from hamilton.execution.grouping import TaskImplementation
+from hamilton.function_modifiers import tag
 
 
 class CMDLineExecutionManager(DefaultExecutionManager):
@@ -54,4 +55,5 @@ def wrapper(*args, **kwargs):
     if inspect.isgeneratorfunction(func):
         # get the return type and set it as the return type of the wrapper
         wrapper.__annotations__["return"] = inspect.signature(func).return_annotation[2]
+    wrapper = tag(cmdline="yes")(wrapper)
     return wrapper
diff --git a/examples/cmdline_orchestrator/funcs.py b/examples/cmdline_orchestrator/funcs.py
@@ -6,21 +6,21 @@
 from hamilton.function_modifiers import tag
 
 
-@tag(cmdline="yes", cache="pickle")
+@tag(cache="pickle")
 @cmdline_decorator
 def echo_1(start: str) -> str:
     time.sleep(2)
     return f'echo "1: {start}"'
 
 
-@tag(cmdline="yes", cache="pickle")
+@tag(cache="pickle")
 @cmdline_decorator
 def echo_2(echo_1: str) -> str:
     time.sleep(2)
     return f'echo "2: {echo_1}"'
 
 
-@tag(cmdline="yes", cache="pickle")
+@tag(cache="pickle")
 @cmdline_decorator
 def echo_2b(echo_1: str) -> [str, CompletedProcess, str]:
     # preprocess
@@ -35,7 +35,7 @@ def echo_2b(echo_1: str) -> [str, CompletedProcess, str]:
     return output
 
 
-@tag(cmdline="yes", cache="pickle")
+@tag(cache="pickle")
 @cmdline_decorator
 def echo_3(echo_2: str, echo_2b: str) -> str:
     return f'echo "3: {echo_2 + ":::" + echo_2b}"'