zrb 1.5.11__py3-none-any.whl → 1.5.13__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.
- zrb/builtin/llm/llm_chat.py +1 -1
- zrb/builtin/llm/tool/__init__.py +0 -0
- zrb/builtin/llm/tool/sub_agent.py +125 -0
- zrb/builtin/llm/tool/web.py +0 -2
- zrb/config.py +0 -3
- zrb/llm_config.py +16 -2
- zrb/task/base_task.py +38 -2
- zrb/task/llm/agent.py +5 -8
- zrb/task/llm/context.py +17 -8
- zrb/task/llm/context_enrichment.py +52 -13
- zrb/task/llm/history_summarization.py +3 -5
- zrb/task/llm/prompt.py +7 -4
- zrb/task/llm/tool_wrapper.py +115 -53
- zrb/task/llm_task.py +16 -1
- zrb/util/attr.py +84 -1
- zrb/util/cli/style.py +147 -0
- zrb/util/cli/subcommand.py +22 -1
- zrb/util/cmd/command.py +18 -0
- zrb/util/cmd/remote.py +15 -0
- zrb/util/codemod/modification_mode.py +4 -0
- zrb/util/codemod/modify_class.py +72 -0
- zrb/util/codemod/modify_class_parent.py +68 -0
- zrb/util/codemod/modify_class_property.py +67 -0
- zrb/util/codemod/modify_dict.py +62 -0
- zrb/util/codemod/modify_function.py +75 -3
- zrb/util/codemod/modify_function_call.py +72 -0
- zrb/util/codemod/modify_method.py +77 -0
- zrb/util/codemod/modify_module.py +10 -0
- zrb/util/cron.py +37 -3
- zrb/util/file.py +32 -0
- zrb/util/git.py +113 -0
- zrb/util/git_subtree.py +58 -0
- zrb/util/group.py +64 -2
- zrb/util/load.py +29 -0
- zrb/util/run.py +9 -0
- zrb/util/string/conversion.py +86 -0
- zrb/util/string/format.py +20 -0
- zrb/util/string/name.py +12 -0
- zrb/util/todo.py +165 -4
- {zrb-1.5.11.dist-info → zrb-1.5.13.dist-info}/METADATA +4 -6
- {zrb-1.5.11.dist-info → zrb-1.5.13.dist-info}/RECORD +43 -41
- {zrb-1.5.11.dist-info → zrb-1.5.13.dist-info}/WHEEL +0 -0
- {zrb-1.5.11.dist-info → zrb-1.5.13.dist-info}/entry_points.txt +0 -0
zrb/util/cmd/remote.py
CHANGED
@@ -10,6 +10,21 @@ def get_remote_cmd_script(
|
|
10
10
|
use_password: bool = False,
|
11
11
|
ssh_key: str = "",
|
12
12
|
) -> str:
|
13
|
+
"""
|
14
|
+
Generate an SSH command script to execute a command on a remote host.
|
15
|
+
|
16
|
+
Args:
|
17
|
+
cmd_script (str): The command script to execute on the remote host.
|
18
|
+
host (str): The remote host address.
|
19
|
+
port (int | str): The SSH port.
|
20
|
+
user (str): The SSH user.
|
21
|
+
password (str): The SSH password (used with sshpass if use_password is True).
|
22
|
+
use_password (bool): Whether to use password authentication with sshpass.
|
23
|
+
ssh_key (str): The path to the SSH private key file.
|
24
|
+
|
25
|
+
Returns:
|
26
|
+
str: The generated SSH command script.
|
27
|
+
"""
|
13
28
|
quoted_script = shlex.quote(cmd_script)
|
14
29
|
if ssh_key != "" and use_password:
|
15
30
|
return f'sshpass -p "{password}" ssh -t -p "{port}" -i "{ssh_key}" "{user}@{host}" {quoted_script}' # noqa
|
zrb/util/codemod/modify_class.py
CHANGED
@@ -4,18 +4,75 @@ from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
|
|
4
4
|
|
5
5
|
|
6
6
|
def replace_class_code(original_code: str, class_name: str, new_code: str) -> str:
|
7
|
+
"""
|
8
|
+
Replace the entire code body of a specified class.
|
9
|
+
|
10
|
+
Args:
|
11
|
+
original_code (str): The original Python code as a string.
|
12
|
+
class_name (str): The name of the class to modify.
|
13
|
+
new_code (str): The new code body for the class as a string.
|
14
|
+
|
15
|
+
Returns:
|
16
|
+
str: The modified Python code as a string.
|
17
|
+
|
18
|
+
Raises:
|
19
|
+
ValueError: If the specified class is not found in the code.
|
20
|
+
"""
|
7
21
|
return _modify_code(original_code, class_name, new_code, REPLACE)
|
8
22
|
|
9
23
|
|
10
24
|
def prepend_code_to_class(original_code: str, class_name: str, new_code: str) -> str:
|
25
|
+
"""
|
26
|
+
Prepend code to the body of a specified class.
|
27
|
+
|
28
|
+
Args:
|
29
|
+
original_code (str): The original Python code as a string.
|
30
|
+
class_name (str): The name of the class to modify.
|
31
|
+
new_code (str): The code to prepend as a string.
|
32
|
+
|
33
|
+
Returns:
|
34
|
+
str: The modified Python code as a string.
|
35
|
+
|
36
|
+
Raises:
|
37
|
+
ValueError: If the specified class is not found in the code.
|
38
|
+
"""
|
11
39
|
return _modify_code(original_code, class_name, new_code, PREPEND)
|
12
40
|
|
13
41
|
|
14
42
|
def append_code_to_class(original_code: str, class_name: str, new_code: str) -> str:
|
43
|
+
"""
|
44
|
+
Append code to the body of a specified class.
|
45
|
+
|
46
|
+
Args:
|
47
|
+
original_code (str): The original Python code as a string.
|
48
|
+
class_name (str): The name of the class to modify.
|
49
|
+
new_code (str): The code to append as a string.
|
50
|
+
|
51
|
+
Returns:
|
52
|
+
str: The modified Python code as a string.
|
53
|
+
|
54
|
+
Raises:
|
55
|
+
ValueError: If the specified class is not found in the code.
|
56
|
+
"""
|
15
57
|
return _modify_code(original_code, class_name, new_code, APPEND)
|
16
58
|
|
17
59
|
|
18
60
|
def _modify_code(original_code: str, class_name: str, new_code: str, mode: int) -> str:
|
61
|
+
"""
|
62
|
+
Modify the code body of a specified class.
|
63
|
+
|
64
|
+
Args:
|
65
|
+
original_code (str): The original Python code as a string.
|
66
|
+
class_name (str): The name of the class to modify.
|
67
|
+
new_code (str): The code to add/replace as a string.
|
68
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
69
|
+
|
70
|
+
Returns:
|
71
|
+
str: The modified Python code as a string.
|
72
|
+
|
73
|
+
Raises:
|
74
|
+
ValueError: If the specified class is not found in the code.
|
75
|
+
"""
|
19
76
|
# Parse the original code into a module
|
20
77
|
module = cst.parse_module(original_code)
|
21
78
|
# Initialize transformer with the class name and method code
|
@@ -30,7 +87,19 @@ def _modify_code(original_code: str, class_name: str, new_code: str, mode: int)
|
|
30
87
|
|
31
88
|
|
32
89
|
class _ClassCodeModifier(cst.CSTTransformer):
|
90
|
+
"""
|
91
|
+
A LibCST transformer to modify the code body of a ClassDef node.
|
92
|
+
"""
|
93
|
+
|
33
94
|
def __init__(self, class_name: str, new_code: str, mode: int):
|
95
|
+
"""
|
96
|
+
Initialize the transformer.
|
97
|
+
|
98
|
+
Args:
|
99
|
+
class_name (str): The name of the target class.
|
100
|
+
new_code (str): The new code body as a string.
|
101
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
102
|
+
"""
|
34
103
|
self.class_name = class_name
|
35
104
|
self.new_code = cst.parse_module(new_code).body
|
36
105
|
self.class_found = False
|
@@ -39,6 +108,9 @@ class _ClassCodeModifier(cst.CSTTransformer):
|
|
39
108
|
def leave_ClassDef(
|
40
109
|
self, original_node: cst.ClassDef, updated_node: cst.ClassDef
|
41
110
|
) -> cst.ClassDef:
|
111
|
+
"""
|
112
|
+
Called when leaving a ClassDef node. Modifies the body of the target class.
|
113
|
+
"""
|
42
114
|
# Check if this is the target class
|
43
115
|
if original_node.name.value == self.class_name:
|
44
116
|
self.class_found = True
|
@@ -6,24 +6,81 @@ from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
|
|
6
6
|
def replace_parent_class(
|
7
7
|
original_code: str, class_name: str, parent_class_name: str
|
8
8
|
) -> str:
|
9
|
+
"""
|
10
|
+
Replace the parent class of a specified class in the provided code.
|
11
|
+
|
12
|
+
Args:
|
13
|
+
original_code (str): The original Python code as a string.
|
14
|
+
class_name (str): The name of the class to modify.
|
15
|
+
parent_class_name (str): The new parent class name.
|
16
|
+
|
17
|
+
Returns:
|
18
|
+
str: The modified Python code as a string.
|
19
|
+
|
20
|
+
Raises:
|
21
|
+
ValueError: If the specified class is not found in the code.
|
22
|
+
"""
|
9
23
|
return _modify_parent_class(original_code, class_name, parent_class_name, REPLACE)
|
10
24
|
|
11
25
|
|
12
26
|
def append_parent_class(
|
13
27
|
original_code: str, class_name: str, parent_class_name: str
|
14
28
|
) -> str:
|
29
|
+
"""
|
30
|
+
Append a parent class to a specified class in the provided code.
|
31
|
+
|
32
|
+
Args:
|
33
|
+
original_code (str): The original Python code as a string.
|
34
|
+
class_name (str): The name of the class to modify.
|
35
|
+
parent_class_name (str): The parent class name to append.
|
36
|
+
|
37
|
+
Returns:
|
38
|
+
str: The modified Python code as a string.
|
39
|
+
|
40
|
+
Raises:
|
41
|
+
ValueError: If the specified class is not found in the code.
|
42
|
+
"""
|
15
43
|
return _modify_parent_class(original_code, class_name, parent_class_name, APPEND)
|
16
44
|
|
17
45
|
|
18
46
|
def prepend_parent_class(
|
19
47
|
original_code: str, class_name: str, parent_class_name: str
|
20
48
|
) -> str:
|
49
|
+
"""
|
50
|
+
Prepend a parent class to a specified class in the provided code.
|
51
|
+
|
52
|
+
Args:
|
53
|
+
original_code (str): The original Python code as a string.
|
54
|
+
class_name (str): The name of the class to modify.
|
55
|
+
parent_class_name (str): The parent class name to prepend.
|
56
|
+
|
57
|
+
Returns:
|
58
|
+
str: The modified Python code as a string.
|
59
|
+
|
60
|
+
Raises:
|
61
|
+
ValueError: If the specified class is not found in the code.
|
62
|
+
"""
|
21
63
|
return _modify_parent_class(original_code, class_name, parent_class_name, PREPEND)
|
22
64
|
|
23
65
|
|
24
66
|
def _modify_parent_class(
|
25
67
|
original_code: str, class_name: str, parent_class_name: str, mode: int
|
26
68
|
) -> str:
|
69
|
+
"""
|
70
|
+
Modify the parent class(es) of a specified class in the provided code.
|
71
|
+
|
72
|
+
Args:
|
73
|
+
original_code (str): The original Python code as a string.
|
74
|
+
class_name (str): The name of the class to modify.
|
75
|
+
parent_class_name (str): The parent class name to add/replace.
|
76
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
77
|
+
|
78
|
+
Returns:
|
79
|
+
str: The modified Python code as a string.
|
80
|
+
|
81
|
+
Raises:
|
82
|
+
ValueError: If the specified class is not found in the code.
|
83
|
+
"""
|
27
84
|
# Parse the original code into a module
|
28
85
|
module = cst.parse_module(original_code)
|
29
86
|
# Initialize transformer with the class name and parent class name
|
@@ -39,6 +96,14 @@ def _modify_parent_class(
|
|
39
96
|
|
40
97
|
class _ParentClassAdder(cst.CSTTransformer):
|
41
98
|
def __init__(self, class_name: str, parent_class_name: str, mode: int):
|
99
|
+
"""
|
100
|
+
Initialize the transformer.
|
101
|
+
|
102
|
+
Args:
|
103
|
+
class_name (str): The name of the target class.
|
104
|
+
parent_class_name (str): The name of the parent class to add/replace.
|
105
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
106
|
+
"""
|
42
107
|
self.class_name = class_name
|
43
108
|
self.parent_class_name = parent_class_name
|
44
109
|
self.class_found = False
|
@@ -47,6 +112,9 @@ class _ParentClassAdder(cst.CSTTransformer):
|
|
47
112
|
def leave_ClassDef(
|
48
113
|
self, original_node: cst.ClassDef, updated_node: cst.ClassDef
|
49
114
|
) -> cst.ClassDef:
|
115
|
+
"""
|
116
|
+
Called when leaving a ClassDef node. Modifies the bases of the target class.
|
117
|
+
"""
|
50
118
|
# Check if this is the target class
|
51
119
|
if original_node.name.value == self.class_name:
|
52
120
|
self.class_found = True
|
@@ -10,6 +10,22 @@ def append_property_to_class(
|
|
10
10
|
annotation: str,
|
11
11
|
default_value: str,
|
12
12
|
) -> str:
|
13
|
+
"""
|
14
|
+
Append a property with type annotation and default value to a specified class.
|
15
|
+
|
16
|
+
Args:
|
17
|
+
original_code (str): The original Python code as a string.
|
18
|
+
class_name (str): The name of the class to modify.
|
19
|
+
property_name (str): The name of the property to add.
|
20
|
+
annotation (str): The type annotation for the property as a string.
|
21
|
+
default_value (str): The default value for the property as a string.
|
22
|
+
|
23
|
+
Returns:
|
24
|
+
str: The modified Python code as a string.
|
25
|
+
|
26
|
+
Raises:
|
27
|
+
ValueError: If the specified class is not found in the code.
|
28
|
+
"""
|
13
29
|
return _modify_class_property(
|
14
30
|
original_code, class_name, property_name, annotation, default_value, APPEND
|
15
31
|
)
|
@@ -22,6 +38,22 @@ def prepend_property_to_class(
|
|
22
38
|
annotation: str,
|
23
39
|
default_value: str,
|
24
40
|
) -> str:
|
41
|
+
"""
|
42
|
+
Prepend a property with type annotation and default value to a specified class.
|
43
|
+
|
44
|
+
Args:
|
45
|
+
original_code (str): The original Python code as a string.
|
46
|
+
class_name (str): The name of the class to modify.
|
47
|
+
property_name (str): The name of the property to add.
|
48
|
+
annotation (str): The type annotation for the property as a string.
|
49
|
+
default_value (str): The default value for the property as a string.
|
50
|
+
|
51
|
+
Returns:
|
52
|
+
str: The modified Python code as a string.
|
53
|
+
|
54
|
+
Raises:
|
55
|
+
ValueError: If the specified class is not found in the code.
|
56
|
+
"""
|
25
57
|
return _modify_class_property(
|
26
58
|
original_code, class_name, property_name, annotation, default_value, PREPEND
|
27
59
|
)
|
@@ -35,6 +67,23 @@ def _modify_class_property(
|
|
35
67
|
default_value: str,
|
36
68
|
mode: int,
|
37
69
|
) -> str:
|
70
|
+
"""
|
71
|
+
Modify a class by adding a property with type annotation and default value.
|
72
|
+
|
73
|
+
Args:
|
74
|
+
original_code (str): The original Python code as a string.
|
75
|
+
class_name (str): The name of the class to modify.
|
76
|
+
property_name (str): The name of the property to add.
|
77
|
+
annotation (str): The type annotation for the property as a string.
|
78
|
+
default_value (str): The default value for the property as a string.
|
79
|
+
mode (int): The modification mode (PREPEND or APPEND).
|
80
|
+
|
81
|
+
Returns:
|
82
|
+
str: The modified Python code as a string.
|
83
|
+
|
84
|
+
Raises:
|
85
|
+
ValueError: If the specified class is not found in the code.
|
86
|
+
"""
|
38
87
|
# Parse the original code into a module
|
39
88
|
module = cst.parse_module(original_code)
|
40
89
|
# Initialize transformer with the class name, property name, annotation, and default value
|
@@ -51,6 +100,11 @@ def _modify_class_property(
|
|
51
100
|
|
52
101
|
|
53
102
|
class _ClassPropertyModifier(cst.CSTTransformer):
|
103
|
+
"""
|
104
|
+
A LibCST transformer to add a property with type annotation and default
|
105
|
+
value to a ClassDef node.
|
106
|
+
"""
|
107
|
+
|
54
108
|
def __init__(
|
55
109
|
self,
|
56
110
|
class_name: str,
|
@@ -59,6 +113,16 @@ class _ClassPropertyModifier(cst.CSTTransformer):
|
|
59
113
|
default_value: str,
|
60
114
|
mode: int,
|
61
115
|
):
|
116
|
+
"""
|
117
|
+
Initialize the transformer.
|
118
|
+
|
119
|
+
Args:
|
120
|
+
class_name (str): The name of the target class.
|
121
|
+
property_name (str): The name of the property to add.
|
122
|
+
annotation (str): The type annotation for the property as a string.
|
123
|
+
default_value (str): The default value for the property as a string.
|
124
|
+
mode (int): The modification mode (PREPEND or APPEND).
|
125
|
+
"""
|
62
126
|
self.class_name = class_name
|
63
127
|
self.property_name = property_name
|
64
128
|
self.annotation = cst.Annotation(cst.parse_expression(annotation))
|
@@ -69,6 +133,9 @@ class _ClassPropertyModifier(cst.CSTTransformer):
|
|
69
133
|
def leave_ClassDef(
|
70
134
|
self, original_node: cst.ClassDef, updated_node: cst.ClassDef
|
71
135
|
) -> cst.ClassDef:
|
136
|
+
"""
|
137
|
+
Called when leaving a ClassDef node. Adds the property to the target class.
|
138
|
+
"""
|
72
139
|
# Check if this is the target class
|
73
140
|
if original_node.name.value == self.class_name:
|
74
141
|
self.class_found = True
|
zrb/util/codemod/modify_dict.py
CHANGED
@@ -6,18 +6,64 @@ from zrb.util.codemod.modification_mode import APPEND, PREPEND
|
|
6
6
|
def prepend_key_to_dict(
|
7
7
|
original_code: str, dictionary_name: str, new_key: str, new_value: str
|
8
8
|
) -> str:
|
9
|
+
"""
|
10
|
+
Prepend a new key-value pair to a dictionary in the provided code.
|
11
|
+
|
12
|
+
Args:
|
13
|
+
original_code (str): The original Python code as a string.
|
14
|
+
dictionary_name (str): The name of the dictionary variable.
|
15
|
+
new_key (str): The key to prepend.
|
16
|
+
new_value (str): The value for the new key.
|
17
|
+
|
18
|
+
Returns:
|
19
|
+
str: The modified Python code as a string.
|
20
|
+
|
21
|
+
Raises:
|
22
|
+
ValueError: If the specified dictionary is not found in the code.
|
23
|
+
"""
|
9
24
|
return _modify_dict(original_code, dictionary_name, new_key, new_value, PREPEND)
|
10
25
|
|
11
26
|
|
12
27
|
def append_key_to_dict(
|
13
28
|
original_code: str, dictionary_name: str, new_key: str, new_value: str
|
14
29
|
) -> str:
|
30
|
+
"""
|
31
|
+
Append a new key-value pair to a dictionary in the provided code.
|
32
|
+
|
33
|
+
Args:
|
34
|
+
original_code (str): The original Python code as a string.
|
35
|
+
dictionary_name (str): The name of the dictionary variable.
|
36
|
+
new_key (str): The key to append.
|
37
|
+
new_value (str): The value for the new key.
|
38
|
+
|
39
|
+
Returns:
|
40
|
+
str: The modified Python code as a string.
|
41
|
+
|
42
|
+
Raises:
|
43
|
+
ValueError: If the specified dictionary is not found in the code.
|
44
|
+
"""
|
15
45
|
return _modify_dict(original_code, dictionary_name, new_key, new_value, APPEND)
|
16
46
|
|
17
47
|
|
18
48
|
def _modify_dict(
|
19
49
|
original_code: str, dictionary_name: str, new_key: str, new_value: str, mode: int
|
20
50
|
) -> str:
|
51
|
+
"""
|
52
|
+
Modify a dictionary in the provided code by adding a new key-value pair.
|
53
|
+
|
54
|
+
Args:
|
55
|
+
original_code (str): The original Python code as a string.
|
56
|
+
dictionary_name (str): The name of the dictionary variable.
|
57
|
+
new_key (str): The key to add.
|
58
|
+
new_value (str): The value for the new key.
|
59
|
+
mode (int): The modification mode (PREPEND or APPEND).
|
60
|
+
|
61
|
+
Returns:
|
62
|
+
str: The modified Python code as a string.
|
63
|
+
|
64
|
+
Raises:
|
65
|
+
ValueError: If the specified dictionary is not found in the code.
|
66
|
+
"""
|
21
67
|
# Parse the original code into a module
|
22
68
|
module = cst.parse_module(original_code)
|
23
69
|
# Initialize the transformer with the necessary information
|
@@ -34,7 +80,20 @@ def _modify_dict(
|
|
34
80
|
|
35
81
|
|
36
82
|
class _DictionaryModifier(cst.CSTTransformer):
|
83
|
+
"""
|
84
|
+
A LibCST transformer to modify a dictionary by adding a new key-value pair.
|
85
|
+
"""
|
86
|
+
|
37
87
|
def __init__(self, dictionary_name: str, new_key: str, new_value: str, mode: int):
|
88
|
+
"""
|
89
|
+
Initialize the transformer.
|
90
|
+
|
91
|
+
Args:
|
92
|
+
dictionary_name (str): The name of the target dictionary variable.
|
93
|
+
new_key (str): The key to add.
|
94
|
+
new_value (str): The value for the new key.
|
95
|
+
mode (int): The modification mode (PREPEND or APPEND).
|
96
|
+
"""
|
38
97
|
self.dictionary_name = dictionary_name
|
39
98
|
self.new_key = new_key
|
40
99
|
self.new_value = new_value
|
@@ -44,6 +103,9 @@ class _DictionaryModifier(cst.CSTTransformer):
|
|
44
103
|
def leave_Assign(
|
45
104
|
self, original_node: cst.Assign, updated_node: cst.Assign
|
46
105
|
) -> cst.Assign:
|
106
|
+
"""
|
107
|
+
Called when leaving an Assign node. Modifies the dictionary if the target matches.
|
108
|
+
"""
|
47
109
|
# Extract the first target from updated_node, which will be an AssignTarget
|
48
110
|
target = updated_node.targets[0]
|
49
111
|
# Check if the target is a Name (which should represent the dictionary)
|
@@ -4,24 +4,81 @@ from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
|
|
4
4
|
|
5
5
|
|
6
6
|
def replace_function_code(original_code: str, function_name: str, new_code: str) -> str:
|
7
|
+
"""
|
8
|
+
Replace the entire code body of a specified function.
|
9
|
+
|
10
|
+
Args:
|
11
|
+
original_code (str): The original Python code as a string.
|
12
|
+
function_name (str): The name of the function to modify.
|
13
|
+
new_code (str): The new code body for the function as a string.
|
14
|
+
|
15
|
+
Returns:
|
16
|
+
str: The modified Python code as a string.
|
17
|
+
|
18
|
+
Raises:
|
19
|
+
ValueError: If the specified function is not found in the code.
|
20
|
+
"""
|
7
21
|
return _modify_function(original_code, function_name, new_code, REPLACE)
|
8
22
|
|
9
23
|
|
10
24
|
def prepend_code_to_function(
|
11
25
|
original_code: str, function_name: str, new_code: str
|
12
26
|
) -> str:
|
27
|
+
"""
|
28
|
+
Prepend code to the body of a specified function.
|
29
|
+
|
30
|
+
Args:
|
31
|
+
original_code (str): The original Python code as a string.
|
32
|
+
function_name (str): The name of the function to modify.
|
33
|
+
new_code (str): The code to prepend as a string.
|
34
|
+
|
35
|
+
Returns:
|
36
|
+
str: The modified Python code as a string.
|
37
|
+
|
38
|
+
Raises:
|
39
|
+
ValueError: If the specified function is not found in the code.
|
40
|
+
"""
|
13
41
|
return _modify_function(original_code, function_name, new_code, PREPEND)
|
14
42
|
|
15
43
|
|
16
44
|
def append_code_to_function(
|
17
45
|
original_code: str, function_name: str, new_code: str
|
18
46
|
) -> str:
|
47
|
+
"""
|
48
|
+
Append code to the body of a specified function.
|
49
|
+
|
50
|
+
Args:
|
51
|
+
original_code (str): The original Python code as a string.
|
52
|
+
function_name (str): The name of the function to modify.
|
53
|
+
new_code (str): The code to append as a string.
|
54
|
+
|
55
|
+
Returns:
|
56
|
+
str: The modified Python code as a string.
|
57
|
+
|
58
|
+
Raises:
|
59
|
+
ValueError: If the specified function is not found in the code.
|
60
|
+
"""
|
19
61
|
return _modify_function(original_code, function_name, new_code, APPEND)
|
20
62
|
|
21
63
|
|
22
64
|
def _modify_function(
|
23
65
|
original_code: str, function_name: str, new_code: str, mode: int
|
24
66
|
) -> str:
|
67
|
+
"""
|
68
|
+
Modify the code body of a specified function.
|
69
|
+
|
70
|
+
Args:
|
71
|
+
original_code (str): The original Python code as a string.
|
72
|
+
function_name (str): The name of the function to modify.
|
73
|
+
new_code (str): The code to add/replace as a string.
|
74
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
75
|
+
|
76
|
+
Returns:
|
77
|
+
str: The modified Python code as a string.
|
78
|
+
|
79
|
+
Raises:
|
80
|
+
ValueError: If the specified function is not found in the code.
|
81
|
+
"""
|
25
82
|
# Parse the original code into a module
|
26
83
|
module = cst.parse_module(original_code)
|
27
84
|
# Initialize the transformer with the necessary information
|
@@ -36,7 +93,19 @@ def _modify_function(
|
|
36
93
|
|
37
94
|
|
38
95
|
class _FunctionCodeModifier(cst.CSTTransformer):
|
96
|
+
"""
|
97
|
+
A LibCST transformer to modify the code body of a FunctionDef node.
|
98
|
+
"""
|
99
|
+
|
39
100
|
def __init__(self, function_name: str, new_code: str, mode: int):
|
101
|
+
"""
|
102
|
+
Initialize the transformer.
|
103
|
+
|
104
|
+
Args:
|
105
|
+
function_name (str): The name of the target function.
|
106
|
+
new_code (str): The new code body as a string.
|
107
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
108
|
+
"""
|
40
109
|
self.function_name = function_name
|
41
110
|
# Use parse_module to handle multiple statements
|
42
111
|
self.new_code = cst.parse_module(new_code).body
|
@@ -44,9 +113,12 @@ class _FunctionCodeModifier(cst.CSTTransformer):
|
|
44
113
|
self.mode = mode
|
45
114
|
|
46
115
|
def leave_FunctionDef(
|
47
|
-
self, original_node: cst.
|
48
|
-
) -> cst.
|
49
|
-
|
116
|
+
self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
|
117
|
+
) -> cst.FunctionDef:
|
118
|
+
"""
|
119
|
+
Called when leaving a FunctionDef node. Modifies the body of the target function.
|
120
|
+
"""
|
121
|
+
# Check if the function matches the target function
|
50
122
|
if original_node.name.value == self.function_name:
|
51
123
|
self.function_found = True
|
52
124
|
if self.mode == REPLACE:
|
@@ -6,24 +6,81 @@ from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
|
|
6
6
|
def replace_function_call_param(
|
7
7
|
original_code: str, function_name: str, new_param: str
|
8
8
|
) -> str:
|
9
|
+
"""
|
10
|
+
Replace the parameters of a specified function call.
|
11
|
+
|
12
|
+
Args:
|
13
|
+
original_code (str): The original Python code as a string.
|
14
|
+
function_name (str): The name of the function call to modify.
|
15
|
+
new_param (str): The new parameter(s) as a string (e.g., "param1, param2=value").
|
16
|
+
|
17
|
+
Returns:
|
18
|
+
str: The modified Python code as a string.
|
19
|
+
|
20
|
+
Raises:
|
21
|
+
ValueError: If the specified function call is not found in the code.
|
22
|
+
"""
|
9
23
|
return _modify_function_call(original_code, function_name, new_param, REPLACE)
|
10
24
|
|
11
25
|
|
12
26
|
def prepend_param_to_function_call(
|
13
27
|
original_code: str, function_name: str, new_param: str
|
14
28
|
) -> str:
|
29
|
+
"""
|
30
|
+
Prepend a parameter to a specified function call.
|
31
|
+
|
32
|
+
Args:
|
33
|
+
original_code (str): The original Python code as a string.
|
34
|
+
function_name (str): The name of the function call to modify.
|
35
|
+
new_param (str): The parameter to prepend as a string.
|
36
|
+
|
37
|
+
Returns:
|
38
|
+
str: The modified Python code as a string.
|
39
|
+
|
40
|
+
Raises:
|
41
|
+
ValueError: If the specified function call is not found in the code.
|
42
|
+
"""
|
15
43
|
return _modify_function_call(original_code, function_name, new_param, PREPEND)
|
16
44
|
|
17
45
|
|
18
46
|
def append_param_to_function_call(
|
19
47
|
original_code: str, function_name: str, new_param: str
|
20
48
|
) -> str:
|
49
|
+
"""
|
50
|
+
Append a parameter to a specified function call.
|
51
|
+
|
52
|
+
Args:
|
53
|
+
original_code (str): The original Python code as a string.
|
54
|
+
function_name (str): The name of the function call to modify.
|
55
|
+
new_param (str): The parameter to append as a string.
|
56
|
+
|
57
|
+
Returns:
|
58
|
+
str: The modified Python code as a string.
|
59
|
+
|
60
|
+
Raises:
|
61
|
+
ValueError: If the specified function call is not found in the code.
|
62
|
+
"""
|
21
63
|
return _modify_function_call(original_code, function_name, new_param, APPEND)
|
22
64
|
|
23
65
|
|
24
66
|
def _modify_function_call(
|
25
67
|
original_code: str, function_name: str, new_param: str, mode: int
|
26
68
|
) -> str:
|
69
|
+
"""
|
70
|
+
Modify the parameters of a specified function call.
|
71
|
+
|
72
|
+
Args:
|
73
|
+
original_code (str): The original Python code as a string.
|
74
|
+
function_name (str): The name of the function call to modify.
|
75
|
+
new_param (str): The parameter(s) to add/replace as a string.
|
76
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
77
|
+
|
78
|
+
Returns:
|
79
|
+
str: The modified Python code as a string.
|
80
|
+
|
81
|
+
Raises:
|
82
|
+
ValueError: If the specified function call is not found in the code.
|
83
|
+
"""
|
27
84
|
# Parse the original code into a module
|
28
85
|
module = cst.parse_module(original_code)
|
29
86
|
# Initialize the transformer with the necessary information
|
@@ -40,7 +97,19 @@ def _modify_function_call(
|
|
40
97
|
|
41
98
|
|
42
99
|
class _FunctionCallParamModifier(cst.CSTTransformer):
|
100
|
+
"""
|
101
|
+
A LibCST transformer to modify the parameters of a function call.
|
102
|
+
"""
|
103
|
+
|
43
104
|
def __init__(self, func_name: str, new_param: str, mode: int):
|
105
|
+
"""
|
106
|
+
Initialize the transformer.
|
107
|
+
|
108
|
+
Args:
|
109
|
+
func_name (str): The name of the target function.
|
110
|
+
new_param (str): The new parameter(s) as a string.
|
111
|
+
mode (int): The modification mode (PREPEND, APPEND, or REPLACE).
|
112
|
+
"""
|
44
113
|
self.func_name = func_name
|
45
114
|
# Parse the new parameter to ensure it’s a valid CST node
|
46
115
|
self.new_param = cst.parse_expression(new_param)
|
@@ -48,6 +117,9 @@ class _FunctionCallParamModifier(cst.CSTTransformer):
|
|
48
117
|
self.mode = mode
|
49
118
|
|
50
119
|
def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
|
120
|
+
"""
|
121
|
+
Called when leaving a Call node. Modifies the arguments of the target function call.
|
122
|
+
"""
|
51
123
|
# Check if the function call name matches the target function
|
52
124
|
if (
|
53
125
|
isinstance(original_node.func, cst.Name)
|