77import sys
88from importlib import import_module
99from pathlib import Path
10- from typing import Any , Dict , List , Optional , Tuple
10+ from typing import Any , Dict , List , Optional
1111
1212import nbformat
1313from nbdev .export import nb_export
@@ -18,8 +18,9 @@ class CodeAnalyzer:
1818 Provides code extraction and transformation functionality
1919
2020 Attributes:
21- script_path (Path): Absolute path to the python script generated.
2221 script_name (str): Name of the generated python script.
22+ script_path (Path): Absolute path to the python script generated.
23+ requirements (List[str]): List of pip libraries found in the script.
2324 exported_script_module (ModuleType): The imported module object of the generated script.
2425 available_modules_in_exported_script (list): List of available attributes in the
2526 exported script.
@@ -43,7 +44,8 @@ def __init__(self, notebook_path: Path, output_path: Path) -> None:
4344 f"{ self .script_name } ,
4445 )
4546 ).resolve ()
46- self .__comment_flow_execution ()
47+ self .requirements = self ._get_requirements ()
48+ self .__modify_experiment_script ()
4749
4850 def __get_exp_name (self , notebook_path : Path ) -> str :
4951 """Extract experiment name from Jupyter notebook
@@ -84,17 +86,70 @@ def __convert_to_python(self, notebook_path: Path, output_path: Path, export_fil
8486
8587 return Path (output_path ).joinpath (export_filename ).resolve ()
8688
87- def __comment_flow_execution (self ) -> None :
88- """Comment out lines containing '.run()' in the specified Python script"""
89+ def __modify_experiment_script (self ) -> None :
90+ """Modifies the given python script by commenting out following code:
91+ - occurences of flflow.run()
92+ - instance of FederatedRuntime
93+ """
94+ runtime_class = "FederatedRuntime"
95+ instantiation_info = self .__extract_class_instantiation_info (runtime_class )
96+ instance_name = instantiation_info .get ("instance_name" , [])
97+
98+ with open (self .script_path , "r" ) as file :
99+ code = "" .join (line for line in file if not line .lstrip ().startswith (("!" , "%" )))
100+
101+ code = self .__comment_flow_execution (code )
102+ code = self .__comment_class_instance (code , instance_name )
103+
104+ with open (self .script_path , "w" ) as file :
105+ file .write (code )
106+
107+ def __comment_class_instance (self , script_code : str , instance_name : List [str ]) -> str :
108+ """
109+ Comments out specified class instance in the provided script
110+ Args:
111+ script_code (str): Script content to be analyzed
112+ instance_name (List[str]): The name of the instance
113+
114+ Returns:
115+ str: The updated script with the specified instance lines commented out
116+ """
117+ tree = ast .parse (script_code )
118+ lines = script_code .splitlines ()
119+ lines_to_comment = set ()
120+ for node in ast .walk (tree ):
121+ if isinstance (node , (ast .Assign , ast .Expr )):
122+ if any (
123+ isinstance (subnode , ast .Name ) and subnode .id in instance_name
124+ for subnode in ast .walk (node )
125+ ):
126+ for i in range (node .lineno - 1 , node .end_lineno ):
127+ lines_to_comment .add (i )
128+ modified_lines = [
129+ f"# { line } if idx in lines_to_comment else line for idx , line in enumerate (lines )
130+ ]
131+ updated_script = "\n " .join (modified_lines )
132+
133+ return updated_script
134+
135+ def __comment_flow_execution (self , script_code : str ) -> str :
136+ """
137+ Comment out lines containing '.run()' in the specified Python script
138+ Args:
139+ script_code(str): Script content to be analyzed
140+
141+ Returns:
142+ str: The modified script with run_statement commented out
143+ """
89144 run_statement = ".run()"
145+ lines = script_code .splitlines ()
146+ for idx , line in enumerate (lines ):
147+ stripped_line = line .strip ()
148+ if not stripped_line .startswith ("#" ) and run_statement in stripped_line :
149+ lines [idx ] = f"# { line }
150+ updated_script = "\n " .join (lines )
90151
91- with self .script_path .open ("r" ) as f :
92- data = f .readlines ()
93- for idx , line in enumerate (data ):
94- if run_statement in line :
95- data [idx ] = f"# { line }
96- with self .script_path .open ("w" ) as f :
97- f .writelines (data )
152+ return updated_script
98153
99154 def __import_generated_script (self ) -> None :
100155 """
@@ -161,26 +216,30 @@ def __get_class_name(self, parent_class) -> Optional[str]:
161216 return attr
162217 raise ValueError ("No flow class found that inherits from FLSpec" )
163218
164- def __extract_class_initializing_args (self , class_name ) -> Dict [str , Any ]:
165- """Provided name of the class returns expected arguments and it's
166- values in form of dictionary.
219+ def __extract_class_instantiation_info (self , class_name : str ) -> Dict [str , Any ]:
220+ """
221+ Extracts the instance name and its initialization arguments (both positional and keyword)
222+ for the given class
167223 Args:
168- class_name (str): The name of the class.
224+ class_name (str): The name of the class
169225
170226 Returns:
171- Dict[str, Any]: A dictionary containing the expected arguments and their values.
227+ Dict[str, Any]: A dictionary containing 'args', 'kwargs', and 'instance_name'
172228 """
173- instantiation_args = {"args" : {}, "kwargs" : {}}
174-
175- with open (self .script_path , "r" ) as s :
176- tree = ast .parse (s .read ())
177-
178- for node in ast .walk (tree ):
179- if isinstance (node , ast .Call ) and isinstance (node .func , ast .Name ):
180- if node .func .id == class_name :
181- # We found an instantiation of the class
182- instantiation_args ["args" ] = self ._extract_positional_args (node .args )
183- instantiation_args ["kwargs" ] = self ._extract_keyword_args (node .keywords )
229+ instantiation_args = {"args" : {}, "kwargs" : {}, "instance_name" : []}
230+
231+ with open (self .script_path , "r" ) as file :
232+ code = "" .join (line for line in file if not line .lstrip ().startswith (("!" , "%" )))
233+ tree = ast .parse (code )
234+ for node in ast .walk (tree ):
235+ if isinstance (node , ast .Assign ) and isinstance (node .value , ast .Call ):
236+ if isinstance (node .value .func , ast .Name ) and node .value .func .id == class_name :
237+ for target in node .targets :
238+ if isinstance (target , ast .Name ):
239+ instantiation_args ["instance_name" ].append (target .id )
240+ # We found an instantiation of the class
241+ instantiation_args ["args" ] = self ._extract_positional_args (node .value .args )
242+ instantiation_args ["kwargs" ] = self ._extract_keyword_args (node .value .keywords )
184243
185244 return instantiation_args
186245
@@ -235,41 +294,25 @@ def _clean_value(self, value: str) -> str:
235294 value = value .lstrip ("[" ).rstrip ("]" )
236295 return value
237296
238- def get_requirements (self ) -> Tuple [ List [str ], List [ int ], List [ str ] ]:
297+ def _get_requirements (self ) -> List [str ]:
239298 """Extract pip libraries from the script
240299
241300 Returns:
242- tuple: A tuple containing:
243- requirements (list of str): List of pip libraries found in the script.
244- line_nos (list of int): List of line numbers where "pip install" commands are found.
245- data (list of str): The entire script data as a list of lines.
301+ requirements (List[str]): List of pip libraries found in the script.
246302 """
247303 data = None
248304 with self .script_path .open ("r" ) as f :
249305 requirements = []
250- line_nos = []
251306 data = f .readlines ()
252- for i , line in enumerate (data ):
307+ for _ , line in enumerate (data ):
253308 line = line .strip ()
254309 if "pip install" in line :
255- line_nos .append (i )
256310 # Avoid commented lines, libraries from *.txt file, or openfl.git
257311 # installation
258312 if not line .startswith ("#" ) and "-r" not in line and "openfl.git" not in line :
259313 requirements .append (f"{ line .split (' ' )[- 1 ].strip ()} \n " )
260314
261- return requirements , line_nos , data
262-
263- def remove_lines (self , data : List [str ], line_nos : List [int ]) -> None :
264- """Removes pip install lines from the script
265- Args:
266- data (List[str]): The entire script data as a list of lines.
267- line_nos (List[int]): List of line numbers where "pip install" commands are found.
268- """
269- with self .script_path .open ("w" ) as f :
270- for i , line in enumerate (data ):
271- if i not in line_nos :
272- f .write (line )
315+ return requirements
273316
274317 def get_flow_class_details (self , parent_class ) -> Dict [str , Any ]:
275318 """
@@ -285,7 +328,7 @@ def get_flow_class_details(self, parent_class) -> Dict[str, Any]:
285328 """
286329 flow_class_name = self .__get_class_name (parent_class )
287330 expected_arguments = self .__get_class_arguments (flow_class_name )
288- init_args = self .__extract_class_initializing_args (flow_class_name )
331+ init_args = self .__extract_class_instantiation_info (flow_class_name )
289332
290333 return {
291334 "flow_class_name" : flow_class_name ,
0 commit comments