stepfunction 0.0.1__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 (28) hide show
  1. {stepfunction-0.0.1 → stepfunction-0.0.3}/LICENSE +1 -1
  2. {stepfunction-0.0.1/src/StepFunction.egg-info → stepfunction-0.0.3}/PKG-INFO +6 -6
  3. {stepfunction-0.0.1 → stepfunction-0.0.3}/README.md +0 -1
  4. {stepfunction-0.0.1 → stepfunction-0.0.3}/pyproject.toml +8 -4
  5. {stepfunction-0.0.1 → stepfunction-0.0.3/src/StepFunction.egg-info}/PKG-INFO +6 -6
  6. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/constants/visualizer.py +1 -1
  7. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/core/step_function/__init__.py +1 -1
  8. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/core/step_function/step_function.py +69 -46
  9. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/core/visualizer/__init__.py +1 -1
  10. stepfunction-0.0.3/src/stepfunction/core/visualizer/visualizer.py +148 -0
  11. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/types/step_types.py +2 -3
  12. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/types/visualizer_types.py +3 -2
  13. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/utils/constants.py +2 -2
  14. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/utils/logger.py +12 -7
  15. stepfunction-0.0.3/src/stepfunction/utils/utils.py +18 -0
  16. stepfunction-0.0.1/src/stepfunction/core/visualizer/visualizer.py +0 -129
  17. stepfunction-0.0.1/src/stepfunction/utils/utils.py +0 -17
  18. {stepfunction-0.0.1 → stepfunction-0.0.3}/setup.cfg +0 -0
  19. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/StepFunction.egg-info/SOURCES.txt +0 -0
  20. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/StepFunction.egg-info/dependency_links.txt +0 -0
  21. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/StepFunction.egg-info/requires.txt +0 -0
  22. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/StepFunction.egg-info/top_level.txt +0 -0
  23. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/constants/__init__.py +0 -0
  24. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/constants/enums.py +0 -0
  25. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/exceptions/__init__.py +0 -0
  26. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/exceptions/step_errors.py +0 -0
  27. {stepfunction-0.0.1 → stepfunction-0.0.3}/src/stepfunction/types/__init__.py +0 -0
  28. {stepfunction-0.0.1 → 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,22 +1,23 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: stepfunction
3
- Version: 0.0.1
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/
9
10
  Keywords: StepFunction,Workflow,Orchestration,Library
10
11
  Classifier: Programming Language :: Python :: 3
11
- Classifier: Programming Language :: Python :: 3.12
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
- Requires-Python: >=3.12
16
+ Requires-Python: >=3.9
17
17
  Description-Content-Type: text/markdown
18
18
  License-File: LICENSE
19
19
  Requires-Dist: graphviz===0.20.3
20
+ Dynamic: license-file
20
21
 
21
22
  # StepFunction
22
23
 
@@ -64,4 +65,3 @@ This project is licensed under the MIT License - see the LICENSE file for detail
64
65
 
65
66
  ## Author
66
67
  Created and maintained by **Vineeth Penugonda**.
67
-
@@ -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.1"
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
- requires-python = ">=3.12"
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
- "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,22 +1,23 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: stepfunction
3
- Version: 0.0.1
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/
9
10
  Keywords: StepFunction,Workflow,Orchestration,Library
10
11
  Classifier: Programming Language :: Python :: 3
11
- Classifier: Programming Language :: Python :: 3.12
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
- Requires-Python: >=3.12
16
+ Requires-Python: >=3.9
17
17
  Description-Content-Type: text/markdown
18
18
  License-File: LICENSE
19
19
  Requires-Dist: graphviz===0.20.3
20
+ Dynamic: license-file
20
21
 
21
22
  # StepFunction
22
23
 
@@ -64,4 +65,3 @@ This project is licensed under the MIT License - see the LICENSE file for detail
64
65
 
65
66
  ## Author
66
67
  Created and maintained by **Vineeth Penugonda**.
67
-
@@ -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,17 +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.core.visualizer import Visualizer
13
- from stepfunction.exceptions.step_errors import (ParallelStepExecutionError,
14
- StepExecutionError)
12
+ from stepfunction.exceptions.step_errors import (
13
+ ParallelStepExecutionError,
14
+ StepExecutionError,
15
+ )
15
16
  from stepfunction.types.step_types import StepParams
16
17
  from stepfunction.utils.logger import setup_logger
17
18
 
@@ -91,7 +92,7 @@ class StepFunction:
91
92
  Parallel Example:
92
93
  step_function.add_step("Step1", func1, parallel=True)
93
94
  step_function.add_step("ParallelStep", {
94
- "task1": func2,
95
+ "task1": func2,
95
96
  "task2": func3
96
97
  }, next_step="Step2", parallel=True)
97
98
 
@@ -103,7 +104,7 @@ class StepFunction:
103
104
 
104
105
  Status Example:
105
106
  status = step_function.status # Will be StepFunctionStatus.INITIALIZED, StepFunctionStatus.RUNNING, StepFunctionStatus.COMPLETED, or StepFunctionStatus.FAILED.
106
- """
107
+ """
107
108
 
108
109
  def __init__(self, name: str):
109
110
  self.__name = name # Name of the step function
@@ -120,7 +121,8 @@ class StepFunction:
120
121
  self.__logger = setup_logger(__name__)
121
122
 
122
123
  self.__logger.debug(
123
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
124
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
125
+ )
124
126
 
125
127
  def add_step(
126
128
  self,
@@ -146,7 +148,13 @@ class StepFunction:
146
148
  "stop_on_failure": stop_on_failure,
147
149
  }
148
150
 
149
- 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
+ ):
150
158
  """Add a sub-step function to the workflow."""
151
159
 
152
160
  if name in self.__steps:
@@ -168,7 +176,7 @@ class StepFunction:
168
176
  "branch": None,
169
177
  "parallel": False,
170
178
  "stop_on_failure": False,
171
- "is_sub_step_function": True
179
+ "is_sub_step_function": True,
172
180
  }
173
181
 
174
182
  def set_start_step(self, name: str):
@@ -184,7 +192,8 @@ class StepFunction:
184
192
  self.__status = StepFunctionStatus.RUNNING
185
193
 
186
194
  self.__logger.debug(
187
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
195
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
196
+ )
188
197
 
189
198
  self.__last_result = initial_input
190
199
 
@@ -193,14 +202,14 @@ class StepFunction:
193
202
  try:
194
203
  if step["parallel"]:
195
204
  results = self._execute_parallel(
196
- step["func"], step["stop_on_failure"])
205
+ step["func"], step["stop_on_failure"]
206
+ )
197
207
 
198
208
  self.__last_result = results
199
209
  self.__context[self.__current_step] = results
200
210
 
201
211
  self.__logger.info(
202
- f"Parallel step '{
203
- self.__current_step}' succeeded with results: {results}"
212
+ f"Parallel step '{self.__current_step}' succeeded with results: {results}"
204
213
  )
205
214
  else:
206
215
  result = await self._execute_step(step["func"], self.__last_result)
@@ -208,14 +217,17 @@ class StepFunction:
208
217
  self.__last_result = result
209
218
  self.__context[self.__current_step] = result
210
219
 
211
- self.__logger.info(
212
- f"Step '{self.__current_step}' succeeded"
213
- )
220
+ self.__logger.info(f"Step '{self.__current_step}' succeeded")
214
221
 
215
- if step['branch'] and result in step['branch']:
216
- self.__current_step = step['branch'][result]
217
- else:
218
- 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"]
219
231
 
220
232
  except Exception as exc:
221
233
  self.__logger.exception(
@@ -226,8 +238,7 @@ class StepFunction:
226
238
 
227
239
  if step["on_failure"]:
228
240
  self.__logger.exception(
229
- f"Executing failure step: {
230
- step['on_failure']} for '{self.__current_step}'"
241
+ f"Executing failure step: {step['on_failure']} for '{self.__current_step}'"
231
242
  )
232
243
 
233
244
  self.__current_step = step["on_failure"]
@@ -236,17 +247,18 @@ class StepFunction:
236
247
  self.__status = StepFunctionStatus.FAILED
237
248
 
238
249
  self.__logger.debug(
239
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
250
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
251
+ )
240
252
  else:
241
253
  self.__logger.exception(
242
- f"No failure step defined for '{
243
- self.__current_step}'. Raising Exception."
254
+ f"No failure step defined for '{self.__current_step}'. Raising Exception."
244
255
  )
245
256
 
246
257
  self.__status = StepFunctionStatus.FAILED
247
258
 
248
259
  self.__logger.debug(
249
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
260
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
261
+ )
250
262
 
251
263
  raise StepExecutionError(exc)
252
264
 
@@ -256,24 +268,34 @@ class StepFunction:
256
268
  self.__status = StepFunctionStatus.COMPLETED
257
269
 
258
270
  self.__logger.debug(
259
- f"StepFunction - {self.__name} - Status - {self.__status.value}")
271
+ f"StepFunction - {self.__name} - Status - {self.__status.value}"
272
+ )
260
273
 
261
274
  async def _execute_step(self, func: Callable, input_value: Any):
262
- """ Execute a single step, handling async functions. """
275
+ """Execute a single step, handling async functions."""
263
276
  if iscoroutinefunction(func):
264
277
  return await func(input_value)
265
278
  else:
266
279
  return func(input_value)
267
280
 
268
- 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
+ ):
269
284
  """Execute the steps in parallel."""
270
285
  results = {}
271
286
  errors = []
272
287
  should_stop_execution = False
273
288
 
289
+ def _run(func, arg):
290
+ if iscoroutinefunction(func):
291
+ return asyncio_run(func(arg))
292
+ return func(arg)
293
+
274
294
  with ThreadPoolExecutor() as executor:
275
- futures = {executor.submit(
276
- 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
+ }
277
299
 
278
300
  for future in as_completed(futures):
279
301
  step_name = futures[future]
@@ -286,7 +308,8 @@ class StepFunction:
286
308
  results[step_name] = result
287
309
  except Exception as exc:
288
310
  self.__logger.exception(
289
- f"Parallel task '{step_name}' failed: {exc}")
311
+ f"Parallel task '{step_name}' failed: {exc}"
312
+ )
290
313
 
291
314
  results[step_name] = exc.args[0]
292
315
 
@@ -294,6 +317,8 @@ class StepFunction:
294
317
 
295
318
  if stop_on_failure:
296
319
  should_stop_execution = True
320
+ for f in futures:
321
+ f.cancel()
297
322
 
298
323
  if errors:
299
324
  self.__logger.error(f"Some parallel tasks failed: {errors}")
@@ -306,6 +331,8 @@ class StepFunction:
306
331
 
307
332
  def visualize(self):
308
333
  """Visualize the workflow."""
334
+ from stepfunction.core.visualizer import Visualizer
335
+
309
336
  visualizer = Visualizer(self.__name, self.__steps)
310
337
 
311
338
  self.__logger.debug("Visualizing the step function")
@@ -318,11 +345,12 @@ class StepFunction:
318
345
 
319
346
  output_file_name = visualizer.output_file_name
320
347
 
321
- self.__logger.debug(f"Rendered the step function to file: {
322
- output_file_name}")
348
+ self.__logger.debug(f"Rendered the step function to file: {output_file_name}")
323
349
 
324
350
  def visualize_to_string(self):
325
351
  """Visualize the workflow as a string."""
352
+ from stepfunction.core.visualizer import Visualizer
353
+
326
354
  visualizer = Visualizer(self.__name, self.__steps)
327
355
 
328
356
  self.__logger.debug("Visualizing the step function")
@@ -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,22 +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
+ from typing import Optional
5
6
 
6
- from stepfunction.utils.constants import (DEFAULT_LOG_LEVEL,
7
- DEFAULT_LOGGING_FORMAT,
8
- ENVIRONMENT_VARIABLE_LOG_LEVEL)
7
+ from stepfunction.utils.constants import (
8
+ DEFAULT_LOG_LEVEL,
9
+ DEFAULT_LOGGING_FORMAT,
10
+ ENVIRONMENT_VARIABLE_LOG_LEVEL,
11
+ )
9
12
  from stepfunction.utils.utils import get_environment_variable
10
13
 
11
14
 
12
- def setup_logger(name: str | None = 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:
13
18
  """
14
19
  Set up and return a logger with the given name.
15
20
  If no name is provided, return the root logger.
16
21
  """
17
22
 
18
23
  LOG_LEVEL = get_environment_variable(
19
- ENVIRONMENT_VARIABLE_LOG_LEVEL, DEFAULT_LOG_LEVEL)
24
+ ENVIRONMENT_VARIABLE_LOG_LEVEL, DEFAULT_LOG_LEVEL
25
+ )
20
26
 
21
27
  logging_config = {
22
28
  "version": 1,
@@ -42,7 +48,6 @@ def setup_logger(name: str | None = None, log_format: str = DEFAULT_LOGGING_FORM
42
48
  "handlers": ["console"],
43
49
  "propagate": False,
44
50
  },
45
-
46
51
  "botocore": {
47
52
  "level": "WARNING", # Set level to WARNING to ignore DEBUG logs
48
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('file_name', f"{self.graph_name}.{
102
- 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)
File without changes