stepfunction 0.0.2__tar.gz → 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. {stepfunction-0.0.2 → stepfunction-0.0.3}/LICENSE +1 -1
  2. {stepfunction-0.0.2/src/StepFunction.egg-info → stepfunction-0.0.3}/PKG-INFO +2 -3
  3. {stepfunction-0.0.2 → stepfunction-0.0.3}/README.md +0 -1
  4. {stepfunction-0.0.2 → stepfunction-0.0.3}/pyproject.toml +6 -2
  5. {stepfunction-0.0.2 → stepfunction-0.0.3/src/StepFunction.egg-info}/PKG-INFO +2 -3
  6. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/StepFunction.egg-info/SOURCES.txt +1 -2
  7. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/constants/visualizer.py +1 -1
  8. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/core/step_function/__init__.py +1 -1
  9. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/core/step_function/step_function.py +62 -39
  10. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/core/visualizer/__init__.py +1 -1
  11. stepfunction-0.0.3/src/stepfunction/core/visualizer/visualizer.py +148 -0
  12. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/types/step_types.py +2 -3
  13. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/types/visualizer_types.py +3 -2
  14. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/utils/constants.py +2 -2
  15. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/utils/logger.py +11 -8
  16. stepfunction-0.0.3/src/stepfunction/utils/utils.py +18 -0
  17. stepfunction-0.0.2/src/stepfunction/core/visualizer/visualizer.py +0 -129
  18. stepfunction-0.0.2/src/stepfunction/utils/utils.py +0 -17
  19. stepfunction-0.0.2/tests/test_car_purchase_workflow.py +0 -103
  20. {stepfunction-0.0.2 → stepfunction-0.0.3}/setup.cfg +0 -0
  21. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/StepFunction.egg-info/dependency_links.txt +0 -0
  22. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/StepFunction.egg-info/requires.txt +0 -0
  23. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/StepFunction.egg-info/top_level.txt +0 -0
  24. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/constants/__init__.py +0 -0
  25. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/constants/enums.py +0 -0
  26. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/exceptions/__init__.py +0 -0
  27. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/exceptions/step_errors.py +0 -0
  28. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/types/__init__.py +0 -0
  29. {stepfunction-0.0.2 → stepfunction-0.0.3}/src/stepfunction/utils/__init__.py +0 -0
@@ -18,4 +18,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
18
  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
19
  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
20
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
- SOFTWARE.
21
+ SOFTWARE.
@@ -1,8 +1,9 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: stepfunction
3
- Version: 0.0.2
3
+ Version: 0.0.3
4
4
  Summary: Step Function Workflow Orchestration Library
5
5
  Author: Vineeth Penugonda
6
+ License-Expression: MIT
6
7
  Project-URL: Homepage, https://github.com/vinecodes/stepfunction
7
8
  Project-URL: Issues, https://github.com/vinecodes/stepfunction/issues
8
9
  Project-URL: Blog_Post, https://blog.vineethp.com/posts/introducingstepfunction/
@@ -11,7 +12,6 @@ Classifier: Programming Language :: Python :: 3
11
12
  Classifier: Programming Language :: Python :: 3.9
12
13
  Classifier: Operating System :: OS Independent
13
14
  Classifier: Development Status :: 3 - Alpha
14
- Classifier: License :: OSI Approved :: MIT License
15
15
  Classifier: Topic :: Software Development :: Libraries
16
16
  Requires-Python: >=3.9
17
17
  Description-Content-Type: text/markdown
@@ -65,4 +65,3 @@ This project is licensed under the MIT License - see the LICENSE file for detail
65
65
 
66
66
  ## Author
67
67
  Created and maintained by **Vineeth Penugonda**.
68
-
@@ -44,4 +44,3 @@ This project is licensed under the MIT License - see the LICENSE file for detail
44
44
 
45
45
  ## Author
46
46
  Created and maintained by **Vineeth Penugonda**.
47
-
@@ -1,18 +1,18 @@
1
1
  [project]
2
2
  name = "stepfunction"
3
- version = "0.0.2"
3
+ version = "0.0.3"
4
4
  authors = [{ name = "Vineeth Penugonda" }]
5
5
  description = "Step Function Workflow Orchestration Library"
6
6
  readme = "README.md"
7
7
  keywords = ["StepFunction", "Workflow", "Orchestration", "Library"]
8
8
  requires-python = ">=3.9"
9
9
  dependencies = ["graphviz === 0.20.3"]
10
+ license = "MIT"
10
11
  classifiers = [
11
12
  "Programming Language :: Python :: 3",
12
13
  "Programming Language :: Python :: 3.9",
13
14
  "Operating System :: OS Independent",
14
15
  "Development Status :: 3 - Alpha",
15
- "License :: OSI Approved :: MIT License",
16
16
  "Topic :: Software Development :: Libraries"
17
17
  ]
18
18
 
@@ -20,6 +20,10 @@ classifiers = [
20
20
  requires = ["setuptools>=61.0"]
21
21
  build-backend = "setuptools.build_meta"
22
22
 
23
+ [tool.ruff.lint]
24
+ select = ["E", "F", "I"]
25
+ ignore = ["E501"]
26
+
23
27
  [tool.setuptools.packages.find]
24
28
  where = ["src"]
25
29
  include = ["stepfunction.*"]
@@ -1,8 +1,9 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: stepfunction
3
- Version: 0.0.2
3
+ Version: 0.0.3
4
4
  Summary: Step Function Workflow Orchestration Library
5
5
  Author: Vineeth Penugonda
6
+ License-Expression: MIT
6
7
  Project-URL: Homepage, https://github.com/vinecodes/stepfunction
7
8
  Project-URL: Issues, https://github.com/vinecodes/stepfunction/issues
8
9
  Project-URL: Blog_Post, https://blog.vineethp.com/posts/introducingstepfunction/
@@ -11,7 +12,6 @@ Classifier: Programming Language :: Python :: 3
11
12
  Classifier: Programming Language :: Python :: 3.9
12
13
  Classifier: Operating System :: OS Independent
13
14
  Classifier: Development Status :: 3 - Alpha
14
- Classifier: License :: OSI Approved :: MIT License
15
15
  Classifier: Topic :: Software Development :: Libraries
16
16
  Requires-Python: >=3.9
17
17
  Description-Content-Type: text/markdown
@@ -65,4 +65,3 @@ This project is licensed under the MIT License - see the LICENSE file for detail
65
65
 
66
66
  ## Author
67
67
  Created and maintained by **Vineeth Penugonda**.
68
-
@@ -26,5 +26,4 @@ src/stepfunction/types/visualizer_types.py
26
26
  src/stepfunction/utils/__init__.py
27
27
  src/stepfunction/utils/constants.py
28
28
  src/stepfunction/utils/logger.py
29
- src/stepfunction/utils/utils.py
30
- tests/test_car_purchase_workflow.py
29
+ src/stepfunction/utils/utils.py
@@ -44,4 +44,4 @@ DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_SHAPE = "boxed"
44
44
  """str: The default node shape for sub-step functions in the visualizer."""
45
45
 
46
46
  DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_STYLE = "dotted"
47
- """ str: The default node style for sub-step functions in the visualizer."""
47
+ """ str: The default node style for sub-step functions in the visualizer."""
@@ -1,3 +1,3 @@
1
1
  from .step_function import StepFunction
2
2
 
3
- __all__ = ["StepFunction"]
3
+ __all__ = ["StepFunction"]
@@ -1,16 +1,18 @@
1
- """ Module to define the StepFunction class.
1
+ """Module to define the StepFunction class.
2
2
 
3
- Author: Vineeth Penugonda
3
+ Author: Vineeth Penugonda
4
4
  """
5
5
 
6
-
6
+ from asyncio import run as asyncio_run
7
7
  from concurrent.futures import ThreadPoolExecutor, as_completed
8
8
  from inspect import iscoroutinefunction
9
9
  from typing import Any, Callable, Dict, Optional, Union
10
10
 
11
11
  from stepfunction.constants.enums import StepFunctionStatus
12
- from stepfunction.exceptions.step_errors import (ParallelStepExecutionError,
13
- StepExecutionError)
12
+ from stepfunction.exceptions.step_errors import (
13
+ ParallelStepExecutionError,
14
+ StepExecutionError,
15
+ )
14
16
  from stepfunction.types.step_types import StepParams
15
17
  from stepfunction.utils.logger import setup_logger
16
18
 
@@ -90,7 +92,7 @@ class StepFunction:
90
92
  Parallel Example:
91
93
  step_function.add_step("Step1", func1, parallel=True)
92
94
  step_function.add_step("ParallelStep", {
93
- "task1": func2,
95
+ "task1": func2,
94
96
  "task2": func3
95
97
  }, next_step="Step2", parallel=True)
96
98
 
@@ -102,7 +104,7 @@ class StepFunction:
102
104
 
103
105
  Status Example:
104
106
  status = step_function.status # Will be StepFunctionStatus.INITIALIZED, StepFunctionStatus.RUNNING, StepFunctionStatus.COMPLETED, or StepFunctionStatus.FAILED.
105
- """
107
+ """
106
108
 
107
109
  def __init__(self, name: str):
108
110
  self.__name = name # Name of the step function
@@ -119,7 +121,8 @@ class StepFunction:
119
121
  self.__logger = setup_logger(__name__)
120
122
 
121
123
  self.__logger.debug(
122
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
124
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
125
+ )
123
126
 
124
127
  def add_step(
125
128
  self,
@@ -145,7 +148,13 @@ class StepFunction:
145
148
  "stop_on_failure": stop_on_failure,
146
149
  }
147
150
 
148
- def add_sub_step_function(self, name: str, sub_step_function: "StepFunction", next_step: Optional[str] = None, on_failure: Optional[str] = None):
151
+ def add_sub_step_function(
152
+ self,
153
+ name: str,
154
+ sub_step_function: "StepFunction",
155
+ next_step: Optional[str] = None,
156
+ on_failure: Optional[str] = None,
157
+ ):
149
158
  """Add a sub-step function to the workflow."""
150
159
 
151
160
  if name in self.__steps:
@@ -167,7 +176,7 @@ class StepFunction:
167
176
  "branch": None,
168
177
  "parallel": False,
169
178
  "stop_on_failure": False,
170
- "is_sub_step_function": True
179
+ "is_sub_step_function": True,
171
180
  }
172
181
 
173
182
  def set_start_step(self, name: str):
@@ -183,7 +192,8 @@ class StepFunction:
183
192
  self.__status = StepFunctionStatus.RUNNING
184
193
 
185
194
  self.__logger.debug(
186
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
195
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
196
+ )
187
197
 
188
198
  self.__last_result = initial_input
189
199
 
@@ -192,7 +202,8 @@ class StepFunction:
192
202
  try:
193
203
  if step["parallel"]:
194
204
  results = self._execute_parallel(
195
- step["func"], step["stop_on_failure"])
205
+ step["func"], step["stop_on_failure"]
206
+ )
196
207
 
197
208
  self.__last_result = results
198
209
  self.__context[self.__current_step] = results
@@ -206,14 +217,17 @@ class StepFunction:
206
217
  self.__last_result = result
207
218
  self.__context[self.__current_step] = result
208
219
 
209
- self.__logger.info(
210
- f"Step '{self.__current_step}' succeeded"
211
- )
220
+ self.__logger.info(f"Step '{self.__current_step}' succeeded")
212
221
 
213
- if step['branch'] and result in step['branch']:
214
- self.__current_step = step['branch'][result]
215
- else:
216
- self.__current_step = step["next_step"]
222
+ next_step = None
223
+
224
+ if step["branch"]:
225
+ if callable(step["branch"]):
226
+ next_step = step["branch"](self.__last_result)
227
+ else:
228
+ next_step = step["branch"].get(self.__last_result)
229
+
230
+ self.__current_step = next_step or step["next_step"]
217
231
 
218
232
  except Exception as exc:
219
233
  self.__logger.exception(
@@ -233,7 +247,8 @@ class StepFunction:
233
247
  self.__status = StepFunctionStatus.FAILED
234
248
 
235
249
  self.__logger.debug(
236
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
250
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
251
+ )
237
252
  else:
238
253
  self.__logger.exception(
239
254
  f"No failure step defined for '{self.__current_step}'. Raising Exception."
@@ -242,7 +257,8 @@ class StepFunction:
242
257
  self.__status = StepFunctionStatus.FAILED
243
258
 
244
259
  self.__logger.debug(
245
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
260
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
261
+ )
246
262
 
247
263
  raise StepExecutionError(exc)
248
264
 
@@ -252,24 +268,34 @@ class StepFunction:
252
268
  self.__status = StepFunctionStatus.COMPLETED
253
269
 
254
270
  self.__logger.debug(
255
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
271
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
272
+ )
256
273
 
257
274
  async def _execute_step(self, func: Callable, input_value: Any):
258
- """ Execute a single step, handling async functions. """
275
+ """Execute a single step, handling async functions."""
259
276
  if iscoroutinefunction(func):
260
277
  return await func(input_value)
261
278
  else:
262
279
  return func(input_value)
263
280
 
264
- def _execute_parallel(self, func_dict: Dict[str, Callable[[Any], Any]], stop_on_failure: bool = False):
281
+ def _execute_parallel(
282
+ self, func_dict: Dict[str, Callable[[Any], Any]], stop_on_failure: bool = False
283
+ ):
265
284
  """Execute the steps in parallel."""
266
285
  results = {}
267
286
  errors = []
268
287
  should_stop_execution = False
269
288
 
289
+ def _run(func, arg):
290
+ if iscoroutinefunction(func):
291
+ return asyncio_run(func(arg))
292
+ return func(arg)
293
+
270
294
  with ThreadPoolExecutor() as executor:
271
- futures = {executor.submit(
272
- func, self.__last_result): step_name for step_name, func in func_dict.items()}
295
+ futures = {
296
+ executor.submit(_run, func, self.__last_result): step_name
297
+ for step_name, func in func_dict.items()
298
+ }
273
299
 
274
300
  for future in as_completed(futures):
275
301
  step_name = futures[future]
@@ -282,7 +308,8 @@ class StepFunction:
282
308
  results[step_name] = result
283
309
  except Exception as exc:
284
310
  self.__logger.exception(
285
- f"Parallel task '{step_name}' failed: {exc}")
311
+ f"Parallel task '{step_name}' failed: {exc}"
312
+ )
286
313
 
287
314
  results[step_name] = exc.args[0]
288
315
 
@@ -290,6 +317,8 @@ class StepFunction:
290
317
 
291
318
  if stop_on_failure:
292
319
  should_stop_execution = True
320
+ for f in futures:
321
+ f.cancel()
293
322
 
294
323
  if errors:
295
324
  self.__logger.error(f"Some parallel tasks failed: {errors}")
@@ -316,8 +345,7 @@ class StepFunction:
316
345
 
317
346
  output_file_name = visualizer.output_file_name
318
347
 
319
- self.__logger.debug(
320
- f"Rendered the step function to file: {output_file_name}")
348
+ self.__logger.debug(f"Rendered the step function to file: {output_file_name}")
321
349
 
322
350
  def visualize_to_string(self):
323
351
  """Visualize the workflow as a string."""
@@ -333,32 +361,27 @@ class StepFunction:
333
361
 
334
362
  @property
335
363
  def name(self):
336
- """ Returns the name of the step function.
337
- """
364
+ """Returns the name of the step function."""
338
365
  return self.__name
339
366
 
340
367
  @property
341
368
  def steps(self):
342
- """ Returns the steps of the step function.
343
- """
369
+ """Returns the steps of the step function."""
344
370
  return self.__steps
345
371
 
346
372
  @property
347
373
  def last_result(self):
348
- """ Returns the result of the last step.
349
- """
374
+ """Returns the result of the last step."""
350
375
  return self.__last_result
351
376
 
352
377
  @property
353
378
  def context(self):
354
- """ Returns the context of the step function.
355
- """
379
+ """Returns the context of the step function."""
356
380
  return self.__context
357
381
 
358
382
  @property
359
383
  def status(self):
360
- """ Returns the status of the step function.
361
- """
384
+ """Returns the status of the step function."""
362
385
  return self.__status
363
386
 
364
387
  def __str__(self):
@@ -1,3 +1,3 @@
1
1
  from .visualizer import Visualizer
2
2
 
3
- __all__ = ["Visualizer"]
3
+ __all__ = ["Visualizer"]
@@ -0,0 +1,148 @@
1
+ """This module contains the visualizer class for the graph model.
2
+
3
+ Author: Vineeth Penugonda
4
+ """
5
+
6
+ from os import getcwd
7
+
8
+ from graphviz import Digraph
9
+
10
+ from stepfunction.constants.visualizer import (
11
+ DEFAULT_VISUALIZER_EXTENSION,
12
+ DEFAULT_VISUALIZER_FAILURE_EDGE_COLOR,
13
+ DEFAULT_VISUALIZER_FAILURE_EDGE_LABEL,
14
+ DEFAULT_VISUALIZER_FOLDER,
15
+ DEFAULT_VISUALIZER_FORMAT,
16
+ DEFAULT_VISUALIZER_PARALLEL_STEP_EDGE_STYLE,
17
+ DEFAULT_VISUALIZER_RENDERER,
18
+ DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_COLOR,
19
+ DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_LABEL,
20
+ DEFAULT_VISUALIZER_STRING_ENCODING,
21
+ DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_SHAPE,
22
+ DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_STYLE,
23
+ DEFAULT_VISUALIZER_SUCCESS_EDGE_LABEL,
24
+ )
25
+ from stepfunction.types.step_types import StepParams
26
+ from stepfunction.types.visualizer_types import RenderStepFunctionParams
27
+
28
+
29
+ class Visualizer:
30
+ """This class is responsible for visualizing the graph model."""
31
+
32
+ def __init__(self, graph_name: str, steps: StepParams = None):
33
+ """Initializes the visualizer."""
34
+
35
+ self.graph_name = graph_name
36
+ self.__steps = steps
37
+
38
+ self.__output_file_name = None
39
+ self.__output_file_path = None
40
+
41
+ self.__dot = Digraph(comment=self.graph_name)
42
+
43
+ def visualize_step_function(self):
44
+ """Visualizes the graph model."""
45
+
46
+ if not self.__steps:
47
+ raise ValueError("No steps found to visualize.")
48
+
49
+ for step_name, step_info in self.__steps.items():
50
+ if (
51
+ "is_sub_step_function" in step_info
52
+ and step_info["is_sub_step_function"]
53
+ ):
54
+ self.__dot.node(
55
+ step_name,
56
+ step_name,
57
+ shape=DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_SHAPE,
58
+ style=DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_STYLE,
59
+ )
60
+ else:
61
+ self.__dot.node(step_name, step_name)
62
+
63
+ if step_info["next_step"]:
64
+ self.__dot.edge(
65
+ step_name,
66
+ step_info["next_step"],
67
+ label=DEFAULT_VISUALIZER_SUCCESS_EDGE_LABEL,
68
+ )
69
+
70
+ if step_info["on_failure"]:
71
+ failure_edge_label = DEFAULT_VISUALIZER_FAILURE_EDGE_LABEL
72
+ failure_edge_color = DEFAULT_VISUALIZER_FAILURE_EDGE_COLOR
73
+
74
+ if step_info.get("stop_on_failure"):
75
+ failure_edge_label = DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_LABEL
76
+ failure_edge_color = DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_COLOR
77
+
78
+ self.__dot.edge(
79
+ step_name,
80
+ step_info["on_failure"],
81
+ label=failure_edge_label,
82
+ color=failure_edge_color,
83
+ )
84
+
85
+ if step_info["parallel"]:
86
+ # func is a dictionary for parallel steps
87
+ parallel_function_names = step_info["func"]
88
+
89
+ with self.__dot.subgraph() as s:
90
+ s.attr(rank="same")
91
+
92
+ for parallel_step_name, func in parallel_function_names.items():
93
+ self.__dot.node(parallel_step_name, parallel_step_name)
94
+ self.__dot.edge(
95
+ step_name,
96
+ parallel_step_name,
97
+ style=DEFAULT_VISUALIZER_PARALLEL_STEP_EDGE_STYLE,
98
+ )
99
+
100
+ if step_info["next_step"]:
101
+ self.__dot.edge(parallel_step_name, step_info["next_step"])
102
+
103
+ if step_info.get("branch"):
104
+ for result, next_step in step_info["branch"].items():
105
+ self.__dot.edge(step_name, next_step, label=f"Branch: {result}")
106
+
107
+ def render_step_function(self, **kwargs: RenderStepFunctionParams):
108
+ """Renders the graph model."""
109
+ current_dir = getcwd()
110
+
111
+ format = kwargs.get("format", DEFAULT_VISUALIZER_FORMAT)
112
+ renderer = kwargs.get("renderer", DEFAULT_VISUALIZER_RENDERER)
113
+
114
+ file_path = kwargs.get(
115
+ "file_path", f"{current_dir}/{DEFAULT_VISUALIZER_FOLDER}"
116
+ )
117
+ file_name = kwargs.get(
118
+ "file_name", f"{self.graph_name}.{DEFAULT_VISUALIZER_EXTENSION}"
119
+ )
120
+
121
+ self.__output_file_path = file_path
122
+ self.__output_file_name = file_name
123
+
124
+ self.__dot.render(
125
+ filename=f"{self.__output_file_path}/{self.__output_file_name}",
126
+ format=format,
127
+ renderer=renderer,
128
+ )
129
+
130
+ def render_step_function_to_string(self, **kwargs: RenderStepFunctionParams):
131
+ """Renders the graph model as a string."""
132
+
133
+ format = kwargs.get("format", DEFAULT_VISUALIZER_FORMAT)
134
+ renderer = kwargs.get("renderer", DEFAULT_VISUALIZER_RENDERER)
135
+
136
+ return self.__dot.pipe(format=format, renderer=renderer).decode(
137
+ DEFAULT_VISUALIZER_STRING_ENCODING
138
+ )
139
+
140
+ @property
141
+ def output_file_name(self):
142
+ """Returns the output file name."""
143
+ return self.__output_file_name
144
+
145
+ @property
146
+ def output_file_path(self):
147
+ """Returns the output file path."""
148
+ return self.__output_file_path
@@ -1,12 +1,11 @@
1
- from typing import Any, Callable, Dict, Optional, TypedDict
1
+ from typing import Any, Callable, Dict, Optional, TypedDict, Union
2
2
 
3
3
 
4
4
  class StepParams(TypedDict, total=False):
5
-
6
5
  func: Callable[[Any], Any]
7
6
  next_step: Optional[str]
8
7
  on_failure: Optional[str]
9
- branch: Optional[Dict[Any, str]]
8
+ branch: Optional[Union[Dict[Any, str], Callable[[Any], Optional[str]]]]
10
9
  parallel: bool
11
10
  stop_on_failure: bool
12
11
  is_sub_step_function: bool
@@ -1,7 +1,8 @@
1
- from typing import TypedDict, Optional
1
+ from typing import Optional, TypedDict
2
+
2
3
 
3
4
  class RenderStepFunctionParams(TypedDict, total=False):
4
5
  file_path: Optional[str]
5
6
  file_name: Optional[str]
6
7
  format: str
7
- renderer: str
8
+ renderer: str
@@ -1,4 +1,4 @@
1
- """ Constants for the project. """
1
+ """Constants for the project."""
2
2
 
3
3
  # Logging
4
4
  DEFAULT_LOG_LEVEL = "INFO"
@@ -9,4 +9,4 @@ DEFAULT_LOGGING_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
9
9
 
10
10
  # Environment variables
11
11
  ENVIRONMENT_VARIABLE_LOG_LEVEL = "LOG_LEVEL"
12
- """ str: The environment variable for the logging level."""
12
+ """ str: The environment variable for the logging level."""
@@ -1,24 +1,28 @@
1
- """ Module for setting up loggers. """
1
+ """Module for setting up loggers."""
2
2
 
3
3
  import logging
4
4
  import logging.config
5
-
6
5
  from typing import Optional
7
6
 
8
- from stepfunction.utils.constants import (DEFAULT_LOG_LEVEL,
9
- DEFAULT_LOGGING_FORMAT,
10
- ENVIRONMENT_VARIABLE_LOG_LEVEL)
7
+ from stepfunction.utils.constants import (
8
+ DEFAULT_LOG_LEVEL,
9
+ DEFAULT_LOGGING_FORMAT,
10
+ ENVIRONMENT_VARIABLE_LOG_LEVEL,
11
+ )
11
12
  from stepfunction.utils.utils import get_environment_variable
12
13
 
13
14
 
14
- def setup_logger(name: Optional[str] = None, log_format: str = DEFAULT_LOGGING_FORMAT) -> logging.Logger:
15
+ def setup_logger(
16
+ name: Optional[str] = None, log_format: str = DEFAULT_LOGGING_FORMAT
17
+ ) -> logging.Logger:
15
18
  """
16
19
  Set up and return a logger with the given name.
17
20
  If no name is provided, return the root logger.
18
21
  """
19
22
 
20
23
  LOG_LEVEL = get_environment_variable(
21
- ENVIRONMENT_VARIABLE_LOG_LEVEL, DEFAULT_LOG_LEVEL)
24
+ ENVIRONMENT_VARIABLE_LOG_LEVEL, DEFAULT_LOG_LEVEL
25
+ )
22
26
 
23
27
  logging_config = {
24
28
  "version": 1,
@@ -44,7 +48,6 @@ def setup_logger(name: Optional[str] = None, log_format: str = DEFAULT_LOGGING_F
44
48
  "handlers": ["console"],
45
49
  "propagate": False,
46
50
  },
47
-
48
51
  "botocore": {
49
52
  "level": "WARNING", # Set level to WARNING to ignore DEBUG logs
50
53
  "handlers": ["console"],
@@ -0,0 +1,18 @@
1
+ """This module contains utility functions for the stepfunction package."""
2
+
3
+ from os import getenv
4
+ from typing import Optional
5
+
6
+
7
+ def get_environment_variable(name: str, default: Optional[str] = None) -> Optional[str]:
8
+ """
9
+ Returns the value of the environment variable with the given name.
10
+
11
+ Args:
12
+ name (str): The name of the environment variable.
13
+ default (str): The default value to return if the environment variable is not set.
14
+
15
+ Returns:
16
+ str: The value of the environment variable.
17
+ """
18
+ return getenv(name, default)
@@ -1,129 +0,0 @@
1
- """ This module contains the visualizer class for the graph model.
2
-
3
- Author: Vineeth Penugonda
4
- """
5
-
6
- from os import getcwd
7
-
8
- from graphviz import Digraph
9
-
10
- from stepfunction.constants.visualizer import (
11
- DEFAULT_VISUALIZER_EXTENSION, DEFAULT_VISUALIZER_FAILURE_EDGE_COLOR,
12
- DEFAULT_VISUALIZER_FAILURE_EDGE_LABEL, DEFAULT_VISUALIZER_FOLDER,
13
- DEFAULT_VISUALIZER_FORMAT, DEFAULT_VISUALIZER_PARALLEL_STEP_EDGE_STYLE,
14
- DEFAULT_VISUALIZER_RENDERER, DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_COLOR,
15
- DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_LABEL,
16
- DEFAULT_VISUALIZER_STRING_ENCODING,
17
- DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_SHAPE,
18
- DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_STYLE,
19
- DEFAULT_VISUALIZER_SUCCESS_EDGE_LABEL)
20
- from stepfunction.types.step_types import StepParams
21
- from stepfunction.types.visualizer_types import RenderStepFunctionParams
22
-
23
-
24
- class Visualizer:
25
- """ This class is responsible for visualizing the graph model.
26
- """
27
-
28
- def __init__(self, graph_name: str, steps: StepParams = None):
29
- """ Initializes the visualizer."""
30
-
31
- self.graph_name = graph_name
32
- self.__steps = steps
33
-
34
- self.__output_file_name = None
35
- self.__output_file_path = None
36
-
37
- self.__dot = Digraph(comment=self.graph_name)
38
-
39
- def visualize_step_function(self):
40
- """ Visualizes the graph model.
41
- """
42
-
43
- if not self.__steps:
44
- raise ValueError("No steps found to visualize.")
45
-
46
- for step_name, step_info in self.__steps.items():
47
-
48
- if 'is_sub_step_function' in step_info and step_info['is_sub_step_function']:
49
- self.__dot.node(step_name, step_name, shape=DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_SHAPE,
50
- style=DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_NODE_STYLE)
51
- else:
52
- self.__dot.node(step_name, step_name)
53
-
54
- if step_info['next_step']:
55
- self.__dot.edge(
56
- step_name, step_info['next_step'], label=DEFAULT_VISUALIZER_SUCCESS_EDGE_LABEL)
57
-
58
- if step_info['on_failure']:
59
- failure_edge_label = DEFAULT_VISUALIZER_FAILURE_EDGE_LABEL
60
- failure_edge_color = DEFAULT_VISUALIZER_FAILURE_EDGE_COLOR
61
-
62
- if step_info.get('stop_on_failure'):
63
- failure_edge_label = DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_LABEL
64
- failure_edge_color = DEFAULT_VISUALIZER_STOP_ON_FAILURE_EDGE_COLOR
65
-
66
- self.__dot.edge(
67
- step_name, step_info['on_failure'], label=failure_edge_label, color=failure_edge_color)
68
-
69
- if step_info['parallel']:
70
-
71
- # func is a dictionary for parallel steps
72
- parallel_function_names = step_info['func']
73
-
74
- with self.__dot.subgraph() as s:
75
- s.attr(rank='same')
76
-
77
- for parallel_step_name, func in parallel_function_names.items():
78
- self.__dot.node(parallel_step_name, parallel_step_name)
79
- self.__dot.edge(
80
- step_name, parallel_step_name, style=DEFAULT_VISUALIZER_PARALLEL_STEP_EDGE_STYLE)
81
-
82
- if step_info['next_step']:
83
- self.__dot.edge(parallel_step_name,
84
- step_info['next_step'])
85
-
86
- if step_info.get('branch'):
87
- for result, next_step in step_info['branch'].items():
88
- self.__dot.edge(step_name, next_step,
89
- label=f"Branch: {result}")
90
-
91
- def render_step_function(self, **kwargs: RenderStepFunctionParams):
92
- """ Renders the graph model.
93
- """
94
- current_dir = getcwd()
95
-
96
- format = kwargs.get('format', DEFAULT_VISUALIZER_FORMAT)
97
- renderer = kwargs.get('renderer', DEFAULT_VISUALIZER_RENDERER)
98
-
99
- file_path = kwargs.get(
100
- 'file_path', f"{current_dir}/{DEFAULT_VISUALIZER_FOLDER}")
101
- file_name = kwargs.get(
102
- 'file_name', f"{self.graph_name}.{DEFAULT_VISUALIZER_EXTENSION}")
103
-
104
- self.__output_file_path = file_path
105
- self.__output_file_name = file_name
106
-
107
- self.__dot.render(
108
- filename=f"{self.__output_file_path}/{self.__output_file_name}", format=format, renderer=renderer)
109
-
110
- def render_step_function_to_string(self, **kwargs: RenderStepFunctionParams):
111
- """ Renders the graph model as a string.
112
- """
113
-
114
- format = kwargs.get('format', DEFAULT_VISUALIZER_FORMAT)
115
- renderer = kwargs.get('renderer', DEFAULT_VISUALIZER_RENDERER)
116
-
117
- return self.__dot.pipe(format=format, renderer=renderer).decode(DEFAULT_VISUALIZER_STRING_ENCODING)
118
-
119
- @property
120
- def output_file_name(self):
121
- """ Returns the output file name.
122
- """
123
- return self.__output_file_name
124
-
125
- @property
126
- def output_file_path(self):
127
- """ Returns the output file path.
128
- """
129
- return self.__output_file_path
@@ -1,17 +0,0 @@
1
- """ This module contains utility functions for the stepfunction package. """
2
-
3
- from os import getenv
4
-
5
-
6
- def get_environment_variable(name: str, default: str = None) -> str:
7
- """
8
- Returns the value of the environment variable with the given name.
9
-
10
- Args:
11
- name (str): The name of the environment variable.
12
- default (str): The default value to return if the environment variable is not set.
13
-
14
- Returns:
15
- str: The value of the environment variable.
16
- """
17
- return getenv(name, default)
@@ -1,103 +0,0 @@
1
- from pathlib import Path
2
- import sys
3
- import types
4
- import unittest
5
-
6
-
7
- ROOT = Path(__file__).resolve().parents[1]
8
- sys.path.insert(0, str(ROOT / "src"))
9
- sys.path.insert(0, str(ROOT))
10
-
11
-
12
- class _FakeSubgraph:
13
- def __init__(self):
14
- self.attrs = {}
15
-
16
- def __enter__(self):
17
- return self
18
-
19
- def __exit__(self, exc_type, exc, tb):
20
- return False
21
-
22
- def attr(self, **kwargs):
23
- self.attrs.update(kwargs)
24
-
25
-
26
- class _FakeDigraph:
27
- def __init__(self, comment=None):
28
- self.comment = comment
29
- self.nodes = []
30
- self.edges = []
31
- self.render_calls = []
32
-
33
- def node(self, *args, **kwargs):
34
- self.nodes.append((args, kwargs))
35
-
36
- def edge(self, *args, **kwargs):
37
- self.edges.append((args, kwargs))
38
-
39
- def subgraph(self):
40
- return _FakeSubgraph()
41
-
42
- def render(self, *args, **kwargs):
43
- self.render_calls.append((args, kwargs))
44
-
45
- def pipe(self, format=None, renderer=None):
46
- return f"digraph {self.comment} ({format}, {renderer})".encode("utf-8")
47
-
48
-
49
- fake_graphviz = types.ModuleType("graphviz")
50
- fake_graphviz.Digraph = _FakeDigraph
51
- sys.modules.setdefault("graphviz", fake_graphviz)
52
-
53
-
54
- from examples.car_purchase_workflow import validate_car_purchase_transaction_workflow
55
- from stepfunction.constants.enums import StepFunctionStatus
56
-
57
-
58
- class CarPurchaseWorkflowTests(unittest.IsolatedAsyncioTestCase):
59
- async def test_valid_car_purchase_workflow_runs_to_completion(self):
60
- workflow = await validate_car_purchase_transaction_workflow(
61
- transaction={
62
- "transaction_id": "CAR-1001",
63
- "customer_email": "buyer@example.com",
64
- "is_valid": True,
65
- },
66
- visualize=True,
67
- )
68
-
69
- self.assertEqual(workflow.status, StepFunctionStatus.COMPLETED)
70
- self.assertEqual(workflow.last_result["workflow"], "done")
71
- self.assertEqual(
72
- workflow.context["Update_Company_DB_Records"]["update_shipping_db_record"],
73
- "shipping_updated",
74
- )
75
- self.assertEqual(
76
- workflow.context["Update_Car_Purchase_Transaction_Status"]["status"],
77
- "completed",
78
- )
79
- self.assertTrue(
80
- workflow.context["Send_Notification_To_User_Email"]["notification_sent"])
81
-
82
- async def test_invalid_car_purchase_skips_parallel_update_steps(self):
83
- workflow = await validate_car_purchase_transaction_workflow(
84
- transaction={
85
- "transaction_id": "CAR-1002",
86
- "customer_email": "buyer@example.com",
87
- "is_valid": False,
88
- },
89
- visualize=False,
90
- )
91
-
92
- self.assertEqual(workflow.status, StepFunctionStatus.COMPLETED)
93
- self.assertEqual(
94
- workflow.context["Check_If_Car_Purchase_Transaction_Is_Valid"], "invalid")
95
- self.assertNotIn("Update_Company_DB_Records", workflow.context)
96
- self.assertEqual(
97
- workflow.context["Send_Notification_To_User_Email"]["status"],
98
- "validation_failed",
99
- )
100
-
101
-
102
- if __name__ == "__main__":
103
- unittest.main()
File without changes