@@ -47,7 +47,7 @@ class CodeAnalyzer(ast.NodeVisitor):
4747 """
4848
4949 def __init__ (self ):
50- self .structure : List [Dict [str , Any ]] = []
50+ self .analyzed_classes : List [Dict [str , Any ]] = []
5151 self .imports : set [str ] = set ()
5252 self .types : set [str ] = set ()
5353 self ._current_class_info : Dict [str , Any ] | None = None
@@ -106,13 +106,19 @@ def _collect_types_from_node(self, node: ast.AST | None) -> None:
106106 if type_str :
107107 self .types .add (type_str )
108108 elif isinstance (node , ast .Subscript ):
109- self ._collect_types_from_node (node .value )
109+ # Add the base type of the subscript (e.g., "List", "Dict")
110+ if isinstance (node .value , ast .Name ):
111+ self .types .add (node .value .id )
112+ self ._collect_types_from_node (node .value ) # Recurse on value just in case
110113 self ._collect_types_from_node (node .slice )
111114 elif isinstance (node , (ast .Tuple , ast .List )):
112115 for elt in node .elts :
113116 self ._collect_types_from_node (elt )
114- elif isinstance (node , ast .Constant ) and isinstance (node .value , str ):
115- self .types .add (node .value )
117+ elif isinstance (node , ast .Constant ):
118+ if isinstance (node .value , str ): # Forward references
119+ self .types .add (node .value )
120+ elif node .value is None : # None type
121+ self .types .add ("None" )
116122 elif isinstance (node , ast .BinOp ) and isinstance (
117123 node .op , ast .BitOr
118124 ): # For | union type
@@ -164,7 +170,7 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None:
164170 type_str = self ._get_type_str (item .annotation )
165171 class_info ["attributes" ].append ({"name" : attr_name , "type" : type_str })
166172
167- self .structure .append (class_info )
173+ self .analyzed_classes .append (class_info )
168174 self ._current_class_info = class_info
169175 self ._depth += 1
170176 self .generic_visit (node )
@@ -260,6 +266,7 @@ def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
260266 # directly within the class body, not inside a method.
261267 elif isinstance (target , ast .Name ) and not self ._is_in_method :
262268 self ._add_attribute (target .id , self ._get_type_str (node .annotation ))
269+ self ._collect_types_from_node (node .annotation )
263270 self .generic_visit (node )
264271
265272
@@ -280,7 +287,7 @@ def parse_code(code: str) -> tuple[List[Dict[str, Any]], set[str], set[str]]:
280287 tree = ast .parse (code )
281288 analyzer = CodeAnalyzer ()
282289 analyzer .visit (tree )
283- return analyzer .structure , analyzer .imports , analyzer .types
290+ return analyzer .analyzed_classes , analyzer .imports , analyzer .types
284291
285292
286293def parse_file (file_path : str ) -> tuple [List [Dict [str , Any ]], set [str ], set [str ]]:
@@ -332,10 +339,10 @@ def list_code_objects(
332339 all_class_keys = []
333340
334341 def process_structure (
335- structure : List [Dict [str , Any ]], file_name : str | None = None
342+ analyzed_classes : List [Dict [str , Any ]], file_name : str | None = None
336343 ):
337344 """Populates the results dictionary from the parsed AST structure."""
338- for class_info in structure :
345+ for class_info in analyzed_classes :
339346 key = class_info ["class_name" ]
340347 if file_name :
341348 key = f"{ key } (in { file_name } )"
@@ -361,13 +368,13 @@ def process_structure(
361368
362369 # Determine if the path is a file or directory and process accordingly
363370 if os .path .isfile (path ) and path .endswith (".py" ):
364- structure , _ , _ = parse_file (path )
365- process_structure (structure )
371+ analyzed_classes , _ , _ = parse_file (path )
372+ process_structure (analyzed_classes )
366373 elif os .path .isdir (path ):
367374 # This assumes `utils.walk_codebase` is defined elsewhere.
368375 for file_path in utils .walk_codebase (path ):
369- structure , _ , _ = parse_file (file_path )
370- process_structure (structure , file_name = os .path .basename (file_path ))
376+ analyzed_classes , _ , _ = parse_file (file_path )
377+ process_structure (analyzed_classes , file_name = os .path .basename (file_path ))
371378
372379 # Return the data in the desired format based on the flags
373380 if not show_methods and not show_attributes :
@@ -419,11 +426,11 @@ def _build_request_arg_schema(
419426 module_name = os .path .splitext (relative_path )[0 ].replace (os .path .sep , "." )
420427
421428 try :
422- structure , _ , _ = parse_file (file_path )
423- if not structure :
429+ analyzed_classes , _ , _ = parse_file (file_path )
430+ if not analyzed_classes :
424431 continue
425432
426- for class_info in structure :
433+ for class_info in analyzed_classes :
427434 class_name = class_info .get ("class_name" , "Unknown" )
428435 if class_name .endswith ("Request" ):
429436 full_class_name = f"{ module_name } .{ class_name } "
@@ -451,11 +458,11 @@ def _process_service_clients(
451458 if "/services/" not in file_path :
452459 continue
453460
454- structure , imports , types = parse_file (file_path )
461+ analyzed_classes , imports , types = parse_file (file_path )
455462 all_imports .update (imports )
456463 all_types .update (types )
457464
458- for class_info in structure :
465+ for class_info in analyzed_classes :
459466 class_name = class_info ["class_name" ]
460467 if not _should_include_class (class_name , class_filters ):
461468 continue
0 commit comments