@@ -187,6 +187,14 @@ def visit_classdef(self, node: nodes.ClassDef) -> None:
187
187
self .compositions_handler .handle (assignattr , node )
188
188
self .handle_assignattr_type (assignattr , node )
189
189
190
+ # Process class attributes
191
+ for local_nodes in node .locals .values ():
192
+ for local_node in local_nodes :
193
+ if isinstance (local_node , nodes .AssignName ) and isinstance (
194
+ local_node .parent , nodes .Assign
195
+ ):
196
+ self .compositions_handler .handle (local_node , node )
197
+
190
198
def visit_functiondef (self , node : nodes .FunctionDef ) -> None :
191
199
"""Visit an astroid.Function node.
192
200
@@ -308,7 +316,9 @@ def set_next(
308
316
pass
309
317
310
318
@abstractmethod
311
- def handle (self , node : nodes .AssignAttr , parent : nodes .ClassDef ) -> None :
319
+ def handle (
320
+ self , node : nodes .AssignAttr | nodes .AssignName , parent : nodes .ClassDef
321
+ ) -> None :
312
322
pass
313
323
314
324
@@ -333,21 +343,29 @@ def set_next(
333
343
return handler
334
344
335
345
@abstractmethod
336
- def handle (self , node : nodes .AssignAttr , parent : nodes .ClassDef ) -> None :
346
+ def handle (
347
+ self , node : nodes .AssignAttr | nodes .AssignName , parent : nodes .ClassDef
348
+ ) -> None :
337
349
if self ._next_handler :
338
350
self ._next_handler .handle (node , parent )
339
351
340
352
341
353
class CompositionsHandler (AbstractRelationshipHandler ):
342
354
"""Handle composition relationships where parent creates child objects."""
343
355
344
- def handle (self , node : nodes .AssignAttr , parent : nodes .ClassDef ) -> None :
356
+ def handle (
357
+ self , node : nodes .AssignAttr | nodes .AssignName , parent : nodes .ClassDef
358
+ ) -> None :
359
+ # If the node is not part of an assignment, pass to next handler
345
360
if not isinstance (node .parent , (nodes .AnnAssign , nodes .Assign )):
346
361
super ().handle (node , parent )
347
362
return
348
363
349
364
value = node .parent .value
350
365
366
+ # Extract the name to handle both AssignAttr and AssignName nodes
367
+ name = node .attrname if isinstance (node , nodes .AssignAttr ) else node .name
368
+
351
369
# Composition: direct object creation (self.x = P())
352
370
if isinstance (value , nodes .Call ):
353
371
inferred_types = utils .infer_node (node )
@@ -356,8 +374,8 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
356
374
# Resolve nodes to actual class definitions
357
375
resolved_types = resolve_to_class_def (element_types )
358
376
359
- current = set (parent .compositions_type [node . attrname ])
360
- parent .compositions_type [node . attrname ] = list (current | resolved_types )
377
+ current = set (parent .compositions_type [name ])
378
+ parent .compositions_type [name ] = list (current | resolved_types )
361
379
return
362
380
363
381
# Composition: comprehensions with object creation (self.x = [P() for ...])
@@ -377,8 +395,8 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
377
395
# Resolve nodes to actual class definitions
378
396
resolved_types = resolve_to_class_def (element_types )
379
397
380
- current = set (parent .compositions_type [node . attrname ])
381
- parent .compositions_type [node . attrname ] = list (current | resolved_types )
398
+ current = set (parent .compositions_type [name ])
399
+ parent .compositions_type [name ] = list (current | resolved_types )
382
400
return
383
401
384
402
# Not a composition, pass to next handler
@@ -388,13 +406,19 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
388
406
class AggregationsHandler (AbstractRelationshipHandler ):
389
407
"""Handle aggregation relationships where parent receives child objects."""
390
408
391
- def handle (self , node : nodes .AssignAttr , parent : nodes .ClassDef ) -> None :
409
+ def handle (
410
+ self , node : nodes .AssignAttr | nodes .AssignName , parent : nodes .ClassDef
411
+ ) -> None :
412
+ # If the node is not part of an assignment, pass to next handler
392
413
if not isinstance (node .parent , (nodes .AnnAssign , nodes .Assign )):
393
414
super ().handle (node , parent )
394
415
return
395
416
396
417
value = node .parent .value
397
418
419
+ # Extract the name to handle both AssignAttr and AssignName nodes
420
+ name = node .attrname if isinstance (node , nodes .AssignAttr ) else node .name
421
+
398
422
# Aggregation: direct assignment (self.x = x)
399
423
if isinstance (value , nodes .Name ):
400
424
inferred_types = utils .infer_node (node )
@@ -403,8 +427,8 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
403
427
# Resolve nodes to actual class definitions
404
428
resolved_types = resolve_to_class_def (element_types )
405
429
406
- current = set (parent .aggregations_type [node . attrname ])
407
- parent .aggregations_type [node . attrname ] = list (current | resolved_types )
430
+ current = set (parent .aggregations_type [name ])
431
+ parent .aggregations_type [name ] = list (current | resolved_types )
408
432
return
409
433
410
434
# Aggregation: comprehensions without object creation (self.x = [existing_obj for ...])
@@ -424,8 +448,8 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
424
448
# Resolve nodes to actual class definitions
425
449
resolved_types = resolve_to_class_def (element_types )
426
450
427
- current = set (parent .aggregations_type [node . attrname ])
428
- parent .aggregations_type [node . attrname ] = list (current | resolved_types )
451
+ current = set (parent .aggregations_type [name ])
452
+ parent .aggregations_type [name ] = list (current | resolved_types )
429
453
return
430
454
431
455
# Not an aggregation, pass to next handler
@@ -435,7 +459,12 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
435
459
class AssociationsHandler (AbstractRelationshipHandler ):
436
460
"""Handle regular association relationships."""
437
461
438
- def handle (self , node : nodes .AssignAttr , parent : nodes .ClassDef ) -> None :
462
+ def handle (
463
+ self , node : nodes .AssignAttr | nodes .AssignName , parent : nodes .ClassDef
464
+ ) -> None :
465
+ # Extract the name to handle both AssignAttr and AssignName nodes
466
+ name = node .attrname if isinstance (node , nodes .AssignAttr ) else node .name
467
+
439
468
# Type annotation only (x: P) -> Association
440
469
# BUT only if there's no actual assignment (to avoid duplicates)
441
470
if isinstance (node .parent , nodes .AnnAssign ) and node .parent .value is None :
@@ -445,18 +474,18 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
445
474
# Resolve nodes to actual class definitions
446
475
resolved_types = resolve_to_class_def (element_types )
447
476
448
- current = set (parent .associations_type [node . attrname ])
449
- parent .associations_type [node . attrname ] = list (current | resolved_types )
477
+ current = set (parent .associations_type [name ])
478
+ parent .associations_type [name ] = list (current | resolved_types )
450
479
return
451
480
452
481
# Everything else is also association (fallback)
453
- current = set (parent .associations_type [node . attrname ])
482
+ current = set (parent .associations_type [name ])
454
483
inferred_types = utils .infer_node (node )
455
484
element_types = extract_element_types (inferred_types )
456
485
457
486
# Resolve Name nodes to actual class definitions
458
487
resolved_types = resolve_to_class_def (element_types )
459
- parent .associations_type [node . attrname ] = list (current | resolved_types )
488
+ parent .associations_type [name ] = list (current | resolved_types )
460
489
461
490
462
491
def resolve_to_class_def (types : set [nodes .NodeNG ]) -> set [nodes .ClassDef ]:
0 commit comments