@@ -511,6 +511,7 @@ class Parser(object):
511
511
def handle_node (self , node , * args , ** kwargs ):
512
512
node_handler = {
513
513
redbaron .nodes .DefNode : self .handle_function ,
514
+ redbaron .nodes .ClassNode : self .handle_class ,
514
515
#astroid.ClassDef: self.handle_class,
515
516
#astroid.ImportFrom: self.handle_from_import,
516
517
}
@@ -548,9 +549,8 @@ def parse(self, filelike, filename):
548
549
549
550
return module
550
551
551
- def handle_function (self , node , nested = False , method = False ):
552
- if nested :
553
- assert not method
552
+ def handle_function (self , node , nested = False , method = False , * args , ** kwargs ):
553
+ if nested and not method :
554
554
cls = NestedFunction
555
555
else :
556
556
cls = Method if method else Function
@@ -596,6 +596,50 @@ def handle_function(self, node, nested=False, method=False):
596
596
597
597
return function
598
598
599
+ def handle_class (self , node , nested = False , * args , ** kwargs ):
600
+ cls = NestedClass if nested else Class
601
+
602
+ docstring = None
603
+
604
+ skip_nodes = (
605
+ redbaron .nodes .EndlNode ,
606
+ redbaron .nodes .CommentNode ,
607
+ )
608
+ for child in node :
609
+ if not isinstance (child , skip_nodes ):
610
+ break
611
+
612
+ if isinstance (child , redbaron .nodes .StringNode ):
613
+ # docstring!
614
+ docstring = child .value
615
+
616
+ children = []
617
+ for child in node :
618
+ result = self .handle_node (child , nested = True , method = True )
619
+ if result is not None :
620
+ children .append (result )
621
+
622
+ start = node .absolute_bounding_box .top_left .line
623
+ bottom_right = node .absolute_bounding_box .bottom_right
624
+ if bottom_right .line == start :
625
+ end = start
626
+ else :
627
+ end = bottom_right .line - 1
628
+
629
+ klass = cls (node .name ,
630
+ self .source ,
631
+ start ,
632
+ end ,
633
+ [decorator .name for decorator in node .decorators ],
634
+ docstring ,
635
+ children ,
636
+ None )
637
+
638
+ for child in klass .children :
639
+ child .parent = klass
640
+
641
+ return klass
642
+
599
643
# TODO: remove
600
644
def __call__ (self , * args , ** kwargs ):
601
645
return self .parse (* args , ** kwargs )
0 commit comments