@@ -87,3 +87,175 @@ def test_import_extraction(self, code_snippet, expected_imports):
8787 expected = sorted (expected_imports )
8888
8989 assert extracted == expected
90+
91+
92+ # --- Tests CodeAnalyzer handling of Attributes ---
93+
94+
95+ class TestCodeAnalyzerAttributes :
96+
97+ @pytest .mark .parametrize (
98+ "code_snippet, expected_structure" ,
99+ [
100+ pytest .param (
101+ """
102+ class MyClass:
103+ CLASS_VAR = 123
104+ """ ,
105+ [
106+ {
107+ "class_name" : "MyClass" ,
108+ "methods" : [],
109+ "attributes" : [{"name" : "CLASS_VAR" , "type" : None }],
110+ }
111+ ],
112+ id = "class_var_assign" ,
113+ ),
114+ pytest .param (
115+ """
116+ class MyClass:
117+ class_var: int = 456
118+ """ ,
119+ [
120+ {
121+ "class_name" : "MyClass" ,
122+ "methods" : [],
123+ "attributes" : [{"name" : "class_var" , "type" : "int" }],
124+ }
125+ ],
126+ id = "class_var_annassign" ,
127+ ),
128+ pytest .param (
129+ """
130+ class MyClass:
131+ class_var: int
132+ """ ,
133+ [
134+ {
135+ "class_name" : "MyClass" ,
136+ "methods" : [],
137+ "attributes" : [{"name" : "class_var" , "type" : "int" }],
138+ }
139+ ],
140+ id = "class_var_annassign_no_value" ,
141+ ),
142+ pytest .param (
143+ """
144+ class MyClass:
145+ def __init__(self):
146+ self.instance_var = 789
147+ """ ,
148+ [
149+ {
150+ "class_name" : "MyClass" ,
151+ "methods" : [
152+ {
153+ "method_name" : "__init__" ,
154+ "args" : [{"name" : "self" , "type" : None }],
155+ "return_type" : None ,
156+ }
157+ ],
158+ "attributes" : [{"name" : "instance_var" , "type" : None }],
159+ }
160+ ],
161+ id = "instance_var_assign" ,
162+ ),
163+ pytest .param (
164+ """
165+ class MyClass:
166+ def __init__(self):
167+ self.instance_var: str = 'hello'
168+ """ ,
169+ [
170+ {
171+ "class_name" : "MyClass" ,
172+ "methods" : [
173+ {
174+ "method_name" : "__init__" ,
175+ "args" : [{"name" : "self" , "type" : None }],
176+ "return_type" : None ,
177+ }
178+ ],
179+ "attributes" : [{"name" : "instance_var" , "type" : "str" }],
180+ }
181+ ],
182+ id = "instance_var_annassign" ,
183+ ),
184+ pytest .param (
185+ """
186+ class MyClass:
187+ def __init__(self):
188+ self.instance_var: str
189+ """ ,
190+ [
191+ {
192+ "class_name" : "MyClass" ,
193+ "methods" : [
194+ {
195+ "method_name" : "__init__" ,
196+ "args" : [{"name" : "self" , "type" : None }],
197+ "return_type" : None ,
198+ }
199+ ],
200+ "attributes" : [{"name" : "instance_var" , "type" : "str" }],
201+ }
202+ ],
203+ id = "instance_var_annassign_no_value" ,
204+ ),
205+ pytest .param (
206+ """
207+ class MyClass:
208+ VAR_A = 1
209+ var_b: int = 2
210+ def __init__(self):
211+ self.var_c = 3
212+ self.var_d: float = 4.0
213+ """ ,
214+ [
215+ {
216+ "class_name" : "MyClass" ,
217+ "methods" : [
218+ {
219+ "method_name" : "__init__" ,
220+ "args" : [{"name" : "self" , "type" : None }],
221+ "return_type" : None ,
222+ }
223+ ],
224+ "attributes" : [
225+ {"name" : "VAR_A" , "type" : None },
226+ {"name" : "var_b" , "type" : "int" },
227+ {"name" : "var_c" , "type" : None },
228+ {"name" : "var_d" , "type" : "float" },
229+ ],
230+ }
231+ ],
232+ id = "mixed_attributes" ,
233+ ),
234+ pytest .param (
235+ "a = 123 # Module level" ,
236+ [],
237+ id = "module_level_assign" ,
238+ ),
239+ pytest .param (
240+ "b: int = 456 # Module level" ,
241+ [],
242+ id = "module_level_annassign" ,
243+ ),
244+ ],
245+ )
246+ def test_attribute_extraction (self , code_snippet : str , expected_structure : list ):
247+ """Tests the extraction of class and instance attributes."""
248+ analyzer = CodeAnalyzer ()
249+ tree = ast .parse (code_snippet )
250+ analyzer .visit (tree )
251+
252+ extracted = analyzer .structure
253+ # Normalize attributes for order-independent comparison
254+ for item in extracted :
255+ if "attributes" in item :
256+ item ["attributes" ].sort (key = lambda x : x ["name" ])
257+ for item in expected_structure :
258+ if "attributes" in item :
259+ item ["attributes" ].sort (key = lambda x : x ["name" ])
260+
261+ assert extracted == expected_structure
0 commit comments