zrb 1.0.0a1__py3-none-any.whl → 1.0.0a3__py3-none-any.whl

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 (156) hide show
  1. zrb/__init__.py +48 -39
  2. zrb/__main__.py +3 -3
  3. zrb/attr/type.py +2 -1
  4. zrb/builtin/__init__.py +40 -2
  5. zrb/builtin/base64.py +32 -0
  6. zrb/builtin/git.py +156 -0
  7. zrb/builtin/git_subtree.py +88 -0
  8. zrb/builtin/group.py +34 -0
  9. zrb/builtin/llm.py +31 -0
  10. zrb/builtin/md5.py +34 -0
  11. zrb/builtin/project/__init__.py +0 -0
  12. zrb/builtin/project/add/__init__.py +0 -0
  13. zrb/builtin/project/add/fastapp.py +72 -0
  14. zrb/builtin/project/add/fastapp_template/.gitignore +4 -0
  15. zrb/builtin/project/add/fastapp_template/README.md +7 -0
  16. zrb/builtin/project/add/fastapp_template/__init__.py +0 -0
  17. zrb/builtin/project/add/fastapp_template/_zrb/config.py +17 -0
  18. zrb/builtin/project/add/fastapp_template/_zrb/group.py +16 -0
  19. zrb/builtin/project/add/fastapp_template/_zrb/helper.py +97 -0
  20. zrb/builtin/project/add/fastapp_template/_zrb/main.py +132 -0
  21. zrb/builtin/project/add/fastapp_template/_zrb/venv_task.py +22 -0
  22. zrb/builtin/project/add/fastapp_template/common/__init__.py +0 -0
  23. zrb/builtin/project/add/fastapp_template/common/app.py +18 -0
  24. zrb/builtin/project/add/fastapp_template/common/db_engine.py +5 -0
  25. zrb/builtin/project/add/fastapp_template/common/db_repository.py +134 -0
  26. zrb/builtin/project/add/fastapp_template/common/error.py +8 -0
  27. zrb/builtin/project/add/fastapp_template/common/schema.py +5 -0
  28. zrb/builtin/project/add/fastapp_template/common/usecase.py +232 -0
  29. zrb/builtin/project/add/fastapp_template/config.py +29 -0
  30. zrb/builtin/project/add/fastapp_template/main.py +7 -0
  31. zrb/builtin/project/add/fastapp_template/migrate.py +3 -0
  32. zrb/builtin/project/add/fastapp_template/module/__init__.py +0 -0
  33. zrb/builtin/project/add/fastapp_template/module/auth/alembic.ini +117 -0
  34. zrb/builtin/project/add/fastapp_template/module/auth/client/api_client.py +7 -0
  35. zrb/builtin/project/add/fastapp_template/module/auth/client/base_client.py +27 -0
  36. zrb/builtin/project/add/fastapp_template/module/auth/client/direct_client.py +6 -0
  37. zrb/builtin/project/add/fastapp_template/module/auth/client/factory.py +9 -0
  38. zrb/builtin/project/add/fastapp_template/module/auth/migration/README +1 -0
  39. zrb/builtin/project/add/fastapp_template/module/auth/migration/env.py +108 -0
  40. zrb/builtin/project/add/fastapp_template/module/auth/migration/script.py.mako +26 -0
  41. zrb/builtin/project/add/fastapp_template/module/auth/migration/versions/3093c7336477_add_user_table.py +37 -0
  42. zrb/builtin/project/add/fastapp_template/module/auth/migration_metadata.py +6 -0
  43. zrb/builtin/project/add/fastapp_template/module/auth/route.py +22 -0
  44. zrb/builtin/project/add/fastapp_template/module/auth/service/__init__.py +0 -0
  45. zrb/builtin/project/add/fastapp_template/module/auth/service/user/__init__.py +0 -0
  46. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/__init__.py +0 -0
  47. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/db_repository.py +39 -0
  48. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/factory.py +13 -0
  49. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/repository.py +34 -0
  50. zrb/builtin/project/add/fastapp_template/module/auth/service/user/usecase.py +45 -0
  51. zrb/builtin/project/add/fastapp_template/module/gateway/alembic.ini +117 -0
  52. zrb/builtin/project/add/fastapp_template/module/gateway/migration/README +1 -0
  53. zrb/builtin/project/add/fastapp_template/module/gateway/migration/env.py +108 -0
  54. zrb/builtin/project/add/fastapp_template/module/gateway/migration/script.py.mako +26 -0
  55. zrb/builtin/project/add/fastapp_template/module/gateway/migration/versions/.gitkeep +0 -0
  56. zrb/builtin/project/add/fastapp_template/module/gateway/migration_metadata.py +3 -0
  57. zrb/builtin/project/add/fastapp_template/module/gateway/route.py +27 -0
  58. zrb/builtin/project/add/fastapp_template/requirements.txt +6 -0
  59. zrb/builtin/project/add/fastapp_template/schema/__init__.py +0 -0
  60. zrb/builtin/project/add/fastapp_template/schema/role.py +31 -0
  61. zrb/builtin/project/add/fastapp_template/schema/user.py +31 -0
  62. zrb/builtin/project/add/fastapp_template/template.env +2 -0
  63. zrb/builtin/project/create/__init__.py +0 -0
  64. zrb/builtin/project/create/create.py +41 -0
  65. zrb/builtin/project/create/project-template/README.md +3 -0
  66. zrb/builtin/project/create/project-template/zrb_init.py +7 -0
  67. zrb/builtin/python.py +11 -0
  68. zrb/builtin/shell/__init__.py +0 -5
  69. zrb/builtin/shell/autocomplete/__init__.py +0 -9
  70. zrb/builtin/shell/autocomplete/bash.py +5 -6
  71. zrb/builtin/shell/autocomplete/subcmd.py +7 -8
  72. zrb/builtin/shell/autocomplete/zsh.py +5 -6
  73. zrb/builtin/todo.py +186 -0
  74. zrb/callback/any_callback.py +1 -1
  75. zrb/callback/callback.py +5 -5
  76. zrb/cmd/cmd_val.py +2 -2
  77. zrb/config.py +4 -1
  78. zrb/content_transformer/any_content_transformer.py +1 -1
  79. zrb/content_transformer/content_transformer.py +2 -2
  80. zrb/context/any_context.py +5 -1
  81. zrb/context/any_shared_context.py +3 -3
  82. zrb/context/context.py +15 -9
  83. zrb/context/shared_context.py +9 -8
  84. zrb/env/__init__.py +0 -3
  85. zrb/env/any_env.py +2 -2
  86. zrb/env/env.py +4 -5
  87. zrb/env/env_file.py +4 -4
  88. zrb/env/env_map.py +4 -4
  89. zrb/group/__init__.py +0 -3
  90. zrb/group/any_group.py +3 -3
  91. zrb/group/group.py +7 -6
  92. zrb/input/any_input.py +1 -1
  93. zrb/input/base_input.py +4 -4
  94. zrb/input/bool_input.py +5 -5
  95. zrb/input/float_input.py +3 -3
  96. zrb/input/int_input.py +3 -3
  97. zrb/input/option_input.py +51 -0
  98. zrb/input/password_input.py +2 -2
  99. zrb/input/str_input.py +1 -1
  100. zrb/input/text_input.py +12 -10
  101. zrb/runner/cli.py +79 -44
  102. zrb/runner/web_app/group_info_ui/controller.py +7 -8
  103. zrb/runner/web_app/group_info_ui/view.html +2 -2
  104. zrb/runner/web_app/home_page/controller.py +7 -6
  105. zrb/runner/web_app/home_page/view.html +2 -2
  106. zrb/runner/web_app/task_ui/controller.py +13 -13
  107. zrb/runner/web_app/task_ui/partial/common-util.js +37 -0
  108. zrb/runner/web_app/task_ui/partial/main.js +9 -2
  109. zrb/runner/web_app/task_ui/partial/show-existing-session.js +20 -5
  110. zrb/runner/web_app/task_ui/partial/visualize-history.js +1 -41
  111. zrb/runner/web_app/task_ui/view.html +4 -2
  112. zrb/runner/web_server.py +137 -211
  113. zrb/runner/web_util.py +5 -35
  114. zrb/session/any_session.py +13 -7
  115. zrb/session/session.py +80 -41
  116. zrb/session_state_log/session_state_log.py +7 -5
  117. zrb/session_state_logger/any_session_state_logger.py +1 -1
  118. zrb/session_state_logger/default_session_state_logger.py +2 -2
  119. zrb/session_state_logger/file_session_state_logger.py +19 -27
  120. zrb/task/any_task.py +8 -3
  121. zrb/task/base_task.py +47 -33
  122. zrb/task/base_trigger.py +11 -12
  123. zrb/task/cmd_task.py +55 -43
  124. zrb/task/http_check.py +8 -8
  125. zrb/task/llm_task.py +160 -0
  126. zrb/task/make_task.py +9 -9
  127. zrb/task/rsync_task.py +7 -7
  128. zrb/task/scaffolder.py +14 -11
  129. zrb/task/scheduler.py +6 -7
  130. zrb/task/task.py +1 -1
  131. zrb/task/tcp_check.py +8 -8
  132. zrb/util/attr.py +19 -3
  133. zrb/util/cli/style.py +71 -2
  134. zrb/util/cli/subcommand.py +2 -2
  135. zrb/util/codemod/__init__.py +0 -0
  136. zrb/util/codemod/add_code_to_class.py +35 -0
  137. zrb/util/codemod/add_code_to_function.py +36 -0
  138. zrb/util/codemod/add_code_to_method.py +55 -0
  139. zrb/util/codemod/add_key_to_dict.py +51 -0
  140. zrb/util/codemod/add_param_to_function_call.py +39 -0
  141. zrb/util/codemod/add_property_to_class.py +55 -0
  142. zrb/util/git.py +156 -0
  143. zrb/util/git_subtree.py +94 -0
  144. zrb/util/group.py +2 -2
  145. zrb/util/llm/tool.py +63 -0
  146. zrb/util/string/conversion.py +7 -0
  147. zrb/util/todo.py +135 -0
  148. {zrb-1.0.0a1.dist-info → zrb-1.0.0a3.dist-info}/METADATA +11 -7
  149. zrb-1.0.0a3.dist-info/RECORD +194 -0
  150. zrb/builtin/shell/_group.py +0 -9
  151. zrb/builtin/shell/autocomplete/_group.py +0 -6
  152. zrb/runner/web_app/any_request_handler.py +0 -24
  153. zrb/runner/web_server.bak.py +0 -208
  154. zrb-1.0.0a1.dist-info/RECORD +0 -120
  155. {zrb-1.0.0a1.dist-info → zrb-1.0.0a3.dist-info}/WHEEL +0 -0
  156. {zrb-1.0.0a1.dist-info → zrb-1.0.0a3.dist-info}/entry_points.txt +0 -0
zrb/task/make_task.py CHANGED
@@ -1,13 +1,13 @@
1
1
  from collections.abc import Callable
2
2
  from typing import Any
3
3
 
4
- from ..context.any_context import AnyContext
5
- from ..context.any_shared_context import AnySharedContext
6
- from ..env.any_env import AnyEnv
7
- from ..group.any_group import AnyGroup
8
- from ..input.any_input import AnyInput
9
- from .any_task import AnyTask
10
- from .base_task import BaseTask
4
+ from zrb.context.any_context import AnyContext
5
+ from zrb.context.any_shared_context import AnySharedContext
6
+ from zrb.env.any_env import AnyEnv
7
+ from zrb.group.any_group import AnyGroup
8
+ from zrb.input.any_input import AnyInput
9
+ from zrb.task.any_task import AnyTask
10
+ from zrb.task.base_task import BaseTask
11
11
 
12
12
 
13
13
  def make_task(
@@ -16,8 +16,8 @@ def make_task(
16
16
  icon: str | None = None,
17
17
  description: str | None = None,
18
18
  cli_only: bool = False,
19
- input: list[AnyInput] | AnyInput | None = None,
20
- env: list[AnyEnv] | AnyEnv | None = None,
19
+ input: list[AnyInput | None] | AnyInput | None = None,
20
+ env: list[AnyEnv | None] | AnyEnv | None = None,
21
21
  execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
22
22
  retries: int = 2,
23
23
  retry_period: float = 0,
zrb/task/rsync_task.py CHANGED
@@ -1,12 +1,12 @@
1
1
  from collections.abc import Callable
2
2
 
3
- from ..attr.type import IntAttr, StrAttr
4
- from ..context.any_context import AnyContext
5
- from ..env.any_env import AnyEnv
6
- from ..input.any_input import AnyInput
7
- from ..util.attr import get_str_attr
8
- from .any_task import AnyTask
9
- from .cmd_task import CmdTask
3
+ from zrb.attr.type import IntAttr, StrAttr
4
+ from zrb.context.any_context import AnyContext
5
+ from zrb.env.any_env import AnyEnv
6
+ from zrb.input.any_input import AnyInput
7
+ from zrb.task.any_task import AnyTask
8
+ from zrb.task.cmd_task import CmdTask
9
+ from zrb.util.attr import get_str_attr
10
10
 
11
11
 
12
12
  class RsyncTask(CmdTask):
zrb/task/scaffolder.py CHANGED
@@ -2,15 +2,15 @@ import os
2
2
  import shutil
3
3
  from collections.abc import Callable
4
4
 
5
- from ..attr.type import BoolAttr, StrAttr
6
- from ..content_transformer.any_content_transformer import AnyContentTransformer
7
- from ..content_transformer.content_transformer import ContentTransformer
8
- from ..context.any_context import AnyContext
9
- from ..env.any_env import AnyEnv
10
- from ..input.any_input import AnyInput
11
- from ..util.attr import get_str_attr
12
- from .any_task import AnyTask
13
- from .base_task import BaseTask
5
+ from zrb.attr.type import BoolAttr, StrAttr
6
+ from zrb.content_transformer.any_content_transformer import AnyContentTransformer
7
+ from zrb.content_transformer.content_transformer import ContentTransformer
8
+ from zrb.context.any_context import AnyContext
9
+ from zrb.env.any_env import AnyEnv
10
+ from zrb.input.any_input import AnyInput
11
+ from zrb.task.any_task import AnyTask
12
+ from zrb.task.base_task import BaseTask
13
+ from zrb.util.attr import get_str_attr
14
14
 
15
15
  TransformConfig = dict[str, str] | Callable[[AnyContext, str], str]
16
16
 
@@ -96,7 +96,10 @@ class Scaffolder(BaseTask):
96
96
  for file_path in file_paths:
97
97
  for transformer in transformers:
98
98
  if transformer.match(ctx, file_path):
99
- transformer.transform_file(ctx, file_path)
99
+ try:
100
+ transformer.transform_file(ctx, file_path)
101
+ except UnicodeDecodeError:
102
+ pass
100
103
 
101
104
  def _copy_path(self, ctx: AnyContext, source_path: str, destination_path: str):
102
105
  """
@@ -119,7 +122,7 @@ class Scaffolder(BaseTask):
119
122
  else:
120
123
  dest_file = os.path.join(
121
124
  destination_path,
122
- self._transform_path(os.path.basename(ctx, source_path)),
125
+ self._transform_path(ctx, os.path.basename(source_path)),
123
126
  )
124
127
  shutil.copy2(source_path, dest_file)
125
128
  ctx.log_info(f"Copied and renamed {source_path} to {dest_file}")
zrb/task/scheduler.py CHANGED
@@ -9,11 +9,10 @@ from zrb.context.any_shared_context import AnySharedContext
9
9
  from zrb.env.any_env import AnyEnv
10
10
  from zrb.input.any_input import AnyInput
11
11
  from zrb.task.any_task import AnyTask
12
-
13
- from ..util.attr import get_str_attr
14
- from ..util.cron import match_cron
15
- from ..xcom.xcom import Xcom
16
- from .base_trigger import BaseTrigger
12
+ from zrb.task.base_trigger import BaseTrigger
13
+ from zrb.util.attr import get_str_attr
14
+ from zrb.util.cron import match_cron
15
+ from zrb.xcom.xcom import Xcom
17
16
 
18
17
 
19
18
  class Scheduler(BaseTrigger):
@@ -25,8 +24,8 @@ class Scheduler(BaseTrigger):
25
24
  icon: str | None = None,
26
25
  description: str | None = None,
27
26
  cli_only: bool = False,
28
- input: list[AnyInput] | AnyInput | None = None,
29
- env: list[AnyEnv] | AnyEnv | None = None,
27
+ input: list[AnyInput | None] | AnyInput | None = None,
28
+ env: list[AnyEnv | None] | AnyEnv | None = None,
30
29
  schedule: StrAttr = None,
31
30
  execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
32
31
  queue_name: fstring | None = None,
zrb/task/task.py CHANGED
@@ -1,4 +1,4 @@
1
- from .base_task import BaseTask
1
+ from zrb.task.base_task import BaseTask
2
2
 
3
3
 
4
4
  class Task(BaseTask):
zrb/task/tcp_check.py CHANGED
@@ -1,14 +1,14 @@
1
1
  import asyncio
2
2
  from collections.abc import Callable
3
3
 
4
- from ..attr.type import IntAttr, StrAttr
5
- from ..context.any_context import AnyContext
6
- from ..context.context import Context
7
- from ..env.any_env import AnyEnv
8
- from ..input.any_input import AnyInput
9
- from ..util.attr import get_int_attr, get_str_attr
10
- from .any_task import AnyTask
11
- from .base_task import BaseTask
4
+ from zrb.attr.type import IntAttr, StrAttr
5
+ from zrb.context.any_context import AnyContext
6
+ from zrb.context.context import Context
7
+ from zrb.env.any_env import AnyEnv
8
+ from zrb.input.any_input import AnyInput
9
+ from zrb.task.any_task import AnyTask
10
+ from zrb.task.base_task import BaseTask
11
+ from zrb.util.attr import get_int_attr, get_str_attr
12
12
 
13
13
 
14
14
  class TcpCheck(BaseTask):
zrb/util/attr.py CHANGED
@@ -1,8 +1,24 @@
1
1
  from typing import Any
2
2
 
3
- from ..attr.type import AnyAttr, BoolAttr, FloatAttr, IntAttr, StrAttr, StrDictAttr
4
- from ..context.any_shared_context import AnySharedContext
5
- from ..util.string.conversion import to_boolean
3
+ from zrb.attr.type import (
4
+ AnyAttr,
5
+ BoolAttr,
6
+ FloatAttr,
7
+ IntAttr,
8
+ StrAttr,
9
+ StrDictAttr,
10
+ StrListAttr,
11
+ )
12
+ from zrb.context.any_shared_context import AnySharedContext
13
+ from zrb.util.string.conversion import to_boolean
14
+
15
+
16
+ def get_str_list_attr(
17
+ shared_ctx: AnySharedContext, attr: StrListAttr | None, auto_render: bool = True
18
+ ) -> list[str]:
19
+ if callable(attr):
20
+ return attr(shared_ctx)
21
+ return {get_str_attr(shared_ctx, val, "", auto_render) for val in attr}
6
22
 
7
23
 
8
24
  def get_str_dict_attr(
zrb/util/cli/style.py CHANGED
@@ -6,8 +6,33 @@ BLUE = 34
6
6
  MAGENTA = 35
7
7
  CYAN = 36
8
8
  WHITE = 37
9
-
10
- VALID_COLORS = [BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE]
9
+ BRIGHT_BLACK = 90
10
+ BRIGHT_RED = 91
11
+ BRIGHT_GREEN = 92
12
+ BRIGHT_YELLOW = 93
13
+ BRIGHT_BLUE = 94
14
+ BRIGHT_MAGENTA = 95
15
+ BRIGHT_CYAN = 96
16
+ BRIGHT_WHITE = 97
17
+
18
+ VALID_COLORS = [
19
+ BLACK,
20
+ RED,
21
+ GREEN,
22
+ YELLOW,
23
+ BLUE,
24
+ MAGENTA,
25
+ CYAN,
26
+ WHITE,
27
+ BRIGHT_BLACK,
28
+ BRIGHT_RED,
29
+ BRIGHT_GREEN,
30
+ BRIGHT_YELLOW,
31
+ BRIGHT_BLUE,
32
+ BRIGHT_MAGENTA,
33
+ BRIGHT_CYAN,
34
+ BRIGHT_WHITE,
35
+ ]
11
36
 
12
37
  BG_BLACK = 40
13
38
  BG_RED = 41
@@ -17,6 +42,14 @@ BG_BLUE = 44
17
42
  BG_MAGENTA = 45
18
43
  BG_CYAN = 46
19
44
  BG_WHITE = 47
45
+ BG_BRIGHT_BLACK = 100
46
+ BG_BRIGHT_RED = 101
47
+ BG_BRIGHT_GREEN = 102
48
+ BG_BRIGHT_YELLOW = 103
49
+ BG_BRIGHT_BLUE = 104
50
+ BG_BRIGHT_MAGENTA = 105
51
+ BG_BRIGHT_CYAN = 106
52
+ BG_BRIGHT_WHITE = 107
20
53
 
21
54
  VALID_BACKGROUNDS = [
22
55
  BG_BLACK,
@@ -27,6 +60,14 @@ VALID_BACKGROUNDS = [
27
60
  BG_MAGENTA,
28
61
  BG_CYAN,
29
62
  BG_WHITE,
63
+ BG_BRIGHT_BLACK,
64
+ BG_BRIGHT_RED,
65
+ BG_BRIGHT_GREEN,
66
+ BG_BRIGHT_YELLOW,
67
+ BG_BRIGHT_BLUE,
68
+ BG_BRIGHT_MAGENTA,
69
+ BG_BRIGHT_CYAN,
70
+ BG_BRIGHT_WHITE,
30
71
  ]
31
72
 
32
73
  BOLD = 1 # Bold or increased intensity
@@ -101,6 +142,34 @@ def stylize_section_header(text: str):
101
142
  return stylize(f" {text} ", color=BLACK, background=BG_WHITE, style=UNDERLINE)
102
143
 
103
144
 
145
+ def stylize_green(text: str):
146
+ return stylize(text, color=GREEN)
147
+
148
+
149
+ def stylize_blue(text: str):
150
+ return stylize(text, color=BLUE)
151
+
152
+
153
+ def stylize_cyan(text: str):
154
+ return stylize(text, color=CYAN)
155
+
156
+
157
+ def stylize_magenta(text: str):
158
+ return stylize(text, color=MAGENTA)
159
+
160
+
161
+ def stylize_yellow(text: str):
162
+ return stylize(text, color=YELLOW)
163
+
164
+
165
+ def stylize_red(text: str):
166
+ return stylize(text, color=RED)
167
+
168
+
169
+ def stylize_bold_green(text: str):
170
+ return stylize(text, color=GREEN, style=BOLD)
171
+
172
+
104
173
  def stylize_bold_yellow(text: str):
105
174
  return stylize(text, color=YELLOW, style=BOLD)
106
175
 
@@ -1,5 +1,5 @@
1
- from ...group.any_group import AnyGroup
2
- from ..group import get_non_empty_subgroups, get_subtasks
1
+ from zrb.group.any_group import AnyGroup
2
+ from zrb.util.group import get_non_empty_subgroups, get_subtasks
3
3
 
4
4
 
5
5
  class SubCommand:
File without changes
@@ -0,0 +1,35 @@
1
+ import libcst as cst
2
+
3
+
4
+ class ClassCodeAdder(cst.CSTTransformer):
5
+ def __init__(self, class_name: str, new_code: str):
6
+ self.class_name = class_name
7
+ self.new_code = cst.parse_statement(new_code)
8
+ self.class_found = False
9
+
10
+ def leave_ClassDef(
11
+ self, original_node: cst.ClassDef, updated_node: cst.ClassDef
12
+ ) -> cst.ClassDef:
13
+ # Check if this is the target class
14
+ if original_node.name.value == self.class_name:
15
+ self.class_found = True
16
+ # Add the method to the class body
17
+ new_body = updated_node.body.with_changes(
18
+ body=updated_node.body.body + (self.new_code,)
19
+ )
20
+ return updated_node.with_changes(body=new_body)
21
+ return updated_node
22
+
23
+
24
+ def add_code_to_class(original_code: str, class_name: str, method_code: str) -> str:
25
+ # Parse the original code into a module
26
+ module = cst.parse_module(original_code)
27
+ # Initialize transformer with the class name and method code
28
+ transformer = ClassCodeAdder(class_name, method_code)
29
+ # Apply the transformation
30
+ modified_module = module.visit(transformer)
31
+ # Check if the class was found
32
+ if not transformer.class_found:
33
+ raise ValueError(f"Class {class_name} not found in the provided code.")
34
+ # Return the modified code
35
+ return modified_module.code
@@ -0,0 +1,36 @@
1
+ import libcst as cst
2
+
3
+
4
+ class FunctionCodeAdder(cst.CSTTransformer):
5
+ def __init__(self, function_name: str, new_code: str):
6
+ self.function_name = function_name
7
+ # Use parse_module to handle multiple statements
8
+ self.new_code = cst.parse_statement(new_code)
9
+ self.function_found = False
10
+
11
+ def leave_FunctionDef(
12
+ self, original_node: cst.ClassDef, updated_node: cst.ClassDef
13
+ ) -> cst.ClassDef:
14
+ # Check if the class matches the target class
15
+ if original_node.name.value == self.function_name:
16
+ self.function_found = True
17
+ # Add the method to the class body
18
+ new_body = updated_node.body.with_changes(
19
+ body=updated_node.body.body + (self.new_code,)
20
+ )
21
+ return updated_node.with_changes(body=new_body)
22
+ return updated_node
23
+
24
+
25
+ def add_code_to_function(original_code: str, function_name: str, new_code: str) -> str:
26
+ # Parse the original code into a module
27
+ module = cst.parse_module(original_code)
28
+ # Initialize the transformer with the necessary information
29
+ transformer = FunctionCodeAdder(function_name, new_code)
30
+ # Apply the transformation
31
+ modified_module = module.visit(transformer)
32
+ # Error handling: raise an error if the class or function is not found
33
+ if not transformer.function_found:
34
+ raise ValueError(f"Function {function_name} not found.")
35
+ # Return the modified code
36
+ return modified_module.code
@@ -0,0 +1,55 @@
1
+ import libcst as cst
2
+
3
+
4
+ class MethodCodeAdder(cst.CSTTransformer):
5
+ def __init__(self, class_name: str, method_name: str, new_code: str):
6
+ self.class_name = class_name
7
+ self.method_name = method_name
8
+ # Use parse_module to handle multiple statements
9
+ self.new_code = cst.parse_statement(new_code)
10
+ self.class_found = False
11
+ self.method_found = False
12
+
13
+ def leave_ClassDef(
14
+ self, original_node: cst.ClassDef, updated_node: cst.ClassDef
15
+ ) -> cst.ClassDef:
16
+ # Check if the class matches the target class
17
+ if original_node.name.value == self.class_name:
18
+ self.class_found = True
19
+ # Now, modify function definitions inside this class
20
+ new_body = []
21
+ for (
22
+ item
23
+ ) in updated_node.body.body: # Access body.body, not just updated_node.body
24
+ if (
25
+ isinstance(item, cst.FunctionDef)
26
+ and item.name.value == self.method_name
27
+ ):
28
+ # Modify the target function by adding the new code
29
+ body_with_new_code = item.body.with_changes(
30
+ body=item.body.body + (self.new_code,) # Add the new code
31
+ )
32
+ new_body.append(item.with_changes(body=body_with_new_code))
33
+ self.method_found = True
34
+ else:
35
+ new_body.append(item)
36
+ return updated_node.with_changes(body=cst.IndentedBlock(new_body))
37
+ return updated_node
38
+
39
+
40
+ def add_code_to_method(
41
+ original_code: str, class_name: str, function_name: str, new_code: str
42
+ ) -> str:
43
+ # Parse the original code into a module
44
+ module = cst.parse_module(original_code)
45
+ # Initialize the transformer with the necessary information
46
+ transformer = MethodCodeAdder(class_name, function_name, new_code)
47
+ # Apply the transformation
48
+ modified_module = module.visit(transformer)
49
+ # Error handling: raise an error if the class or function is not found
50
+ if not transformer.class_found:
51
+ raise ValueError(f"Class {class_name} not found in the provided code.")
52
+ if not transformer.method_found:
53
+ raise ValueError(f"Method {function_name} not found in class {class_name}.")
54
+ # Return the modified code
55
+ return modified_module.code
@@ -0,0 +1,51 @@
1
+ import libcst as cst
2
+
3
+
4
+ class DictionaryModifier(cst.CSTTransformer):
5
+ def __init__(self, dictionary_name: str, new_key: str, new_value: str):
6
+ self.dictionary_name = dictionary_name
7
+ self.new_key = new_key
8
+ self.new_value = new_value
9
+ self.found = False
10
+
11
+ def leave_Assign(
12
+ self, original_node: cst.Assign, updated_node: cst.Assign
13
+ ) -> cst.Assign:
14
+ # Extract the first target from updated_node, which will be an AssignTarget
15
+ target = updated_node.targets[0]
16
+ # Check if the target is a Name (which should represent the dictionary)
17
+ if (
18
+ isinstance(target.target, cst.Name)
19
+ and target.target.value == self.dictionary_name
20
+ ):
21
+ # Check if it's a dictionary initialization (e.g., my_dict = {...})
22
+ if isinstance(updated_node.value, cst.Dict):
23
+ self.found = True
24
+ # Add new key-value pair to the existing dictionary
25
+ new_entries = updated_node.value.elements + (
26
+ cst.DictElement(
27
+ key=cst.SimpleString(f'"{self.new_key}"'),
28
+ value=cst.SimpleString(f'"{self.new_value}"'),
29
+ ),
30
+ )
31
+ new_dict = updated_node.value.with_changes(elements=new_entries)
32
+ return updated_node.with_changes(value=new_dict)
33
+ return updated_node
34
+
35
+
36
+ def add_key_to_dict(
37
+ original_code: str, dictionary_name: str, new_key: str, new_value: str
38
+ ) -> str:
39
+ # Parse the original code into a module
40
+ module = cst.parse_module(original_code)
41
+ # Initialize the transformer with the necessary information
42
+ transformer = DictionaryModifier(dictionary_name, new_key, new_value)
43
+ # Apply the transformation
44
+ modified_module = module.visit(transformer)
45
+ # Error handling: raise an error if the dictionary is not found
46
+ if not transformer.found:
47
+ raise ValueError(
48
+ f"Dictionary {dictionary_name} not found in the provided code."
49
+ )
50
+ # Return the modified code
51
+ return modified_module.code
@@ -0,0 +1,39 @@
1
+ import libcst as cst
2
+
3
+
4
+ class FunctionCallParamAdder(cst.CSTTransformer):
5
+ def __init__(self, func_name: str, new_param: str):
6
+ self.func_name = func_name
7
+ # Parse the new parameter to ensure it’s a valid CST node
8
+ self.new_param = cst.parse_expression(new_param)
9
+ self.param_added = False
10
+
11
+ def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
12
+ # Check if the function call name matches the target function
13
+ if (
14
+ isinstance(original_node.func, cst.Name)
15
+ and original_node.func.value == self.func_name
16
+ ): # noqa
17
+ # Add the new parameter to the function call arguments
18
+ new_args = updated_node.args + (cst.Arg(value=self.new_param),)
19
+ self.param_added = True
20
+ return updated_node.with_changes(args=new_args)
21
+ return updated_node
22
+
23
+
24
+ def add_param_to_function_call(
25
+ original_code: str, func_name: str, new_param: str
26
+ ) -> str:
27
+ # Parse the original code into a module
28
+ module = cst.parse_module(original_code)
29
+ # Initialize the transformer with the necessary information
30
+ transformer = FunctionCallParamAdder(func_name, new_param)
31
+ # Apply the transformation
32
+ modified_module = module.visit(transformer)
33
+ # Error handling: raise an error if the function call is not found
34
+ if not transformer.param_added:
35
+ raise ValueError(
36
+ f"Function call to {func_name} not found in the provided code."
37
+ )
38
+ # Return the modified code
39
+ return modified_module.code
@@ -0,0 +1,55 @@
1
+ import libcst as cst
2
+
3
+
4
+ class ClassPropertyAdder(cst.CSTTransformer):
5
+ def __init__(
6
+ self, class_name: str, property_name: str, annotation: str, default_value: str
7
+ ):
8
+ self.class_name = class_name
9
+ self.property_name = property_name
10
+ self.annotation = cst.Annotation(cst.parse_expression(annotation))
11
+ self.default_value = cst.parse_expression(default_value)
12
+ self.class_found = False
13
+
14
+ def leave_ClassDef(
15
+ self, original_node: cst.ClassDef, updated_node: cst.ClassDef
16
+ ) -> cst.ClassDef:
17
+ # Check if this is the target class
18
+ if original_node.name.value == self.class_name:
19
+ self.class_found = True
20
+ # Create the annotated property with a default value
21
+ new_property = cst.SimpleStatementLine(
22
+ body=[
23
+ cst.AnnAssign(
24
+ target=cst.Name(self.property_name),
25
+ annotation=self.annotation,
26
+ value=self.default_value,
27
+ )
28
+ ]
29
+ )
30
+ # Insert the new property at the start of the class body with a newline
31
+ new_body = cst.IndentedBlock(body=(new_property,) + updated_node.body.body)
32
+ return updated_node.with_changes(body=new_body)
33
+ return updated_node
34
+
35
+
36
+ def add_property_to_class(
37
+ original_code: str,
38
+ class_name: str,
39
+ property_name: str,
40
+ annotation: str,
41
+ default_value: str,
42
+ ) -> str:
43
+ # Parse the original code into a module
44
+ module = cst.parse_module(original_code)
45
+ # Initialize transformer with the class name, property name, annotation, and default value
46
+ transformer = ClassPropertyAdder(
47
+ class_name, property_name, annotation, default_value
48
+ )
49
+ # Apply the transformation
50
+ modified_module = module.visit(transformer)
51
+ # Check if the class was found
52
+ if not transformer.class_found:
53
+ raise ValueError(f"Class {class_name} not found in the provided code.")
54
+ # Return the modified code
55
+ return modified_module.code