Implement methods using AST

SignatureVariablesCollector: remove attr constructor_source from
constructor since it is not used. Add skip_self attr for class
static methods.

Add ClassVisitor and MethodVisitor classes to astvisitors.py
inspectmethod function uses AST to find methods.

Add attributes arguments, is_static, is_class, return_type to
UmlMethod class.

Known limitation: the return type of method is not implemented.

Author: grayfox88

py2puml/domain/umlclass.py

@@ -1,5 +1,5 @@
-from typing import List
-from dataclasses import dataclass
+from typing import List, Dict
+from dataclasses import dataclass, field
 
 from py2puml.domain.umlitem import UmlItem
 
@@ -10,16 +10,23 @@ class UmlAttribute:
     type: str
     static: bool
 
-
 @dataclass
 class UmlMethod:
     name: str
-    signature: str
+    arguments: Dict = field(default_factory=dict)
+    is_static: bool = False
+    is_class: bool = False
+    return_type: str = None
+
+    @property
+    def signature(self):
+        if self.arguments:
+            return ', '.join([f'{arg_type} {arg_name}' for arg_name, arg_type in self.arguments.items()])
+        return ''
 
 
 @dataclass
 class UmlClass(UmlItem):
     attributes: List[UmlAttribute]
-    # TODO move to UmlItem?
     methods: List[UmlMethod]
     is_abstract: bool = False

py2puml/export/puml.py

@@ -12,7 +12,7 @@
 PUML_FILE_END = '@enduml\n'
 PUML_ITEM_START_TPL = '{item_type} {item_fqn} {{\n'
 PUML_ATTR_TPL = '  {attr_name}: {attr_type}{staticity}\n'
-PUML_METHOD_TPL = '  {name}{signature}\n'
+PUML_METHOD_TPL = '  {staticity} {name}({signature})\n'
 PUML_ITEM_END = '}\n'
 PUML_RELATION_TPL = '{source_fqn} {rel_type}-- {target_fqn}\n'
 
@@ -40,7 +40,7 @@ def to_puml_content(diagram_name: str, uml_items: List[UmlItem], uml_relations:
             for uml_attr in uml_class.attributes:
                 yield PUML_ATTR_TPL.format(attr_name=uml_attr.name, attr_type=uml_attr.type, staticity=FEATURE_STATIC if uml_attr.static else FEATURE_INSTANCE)
             for uml_method in uml_class.methods:
-                yield PUML_METHOD_TPL.format(name=uml_method.name, signature=uml_method.signature)
+                yield PUML_METHOD_TPL.format(name=uml_method.name, signature=uml_method.signature, staticity=FEATURE_STATIC if uml_method.is_static else FEATURE_INSTANCE)
             yield PUML_ITEM_END
         else:
             raise TypeError(f'cannot process uml_item of type {uml_item.__class__}')

py2puml/inspection/inspectclass.py

@@ -1,20 +1,20 @@
 from importlib import import_module
-from inspect import isabstract, getmembers, signature
+from inspect import isabstract, getsource
 from re import compile as re_compile
 from typing import Type, List, Dict
-
+from ast import parse, AST
 
 from py2puml.domain.umlitem import UmlItem
 from py2puml.domain.umlclass import UmlClass, UmlAttribute, UmlMethod
 from py2puml.domain.umlrelation import UmlRelation, RelType
-from py2puml.parsing.astvisitors import shorten_compound_type_annotation
+from py2puml.parsing.astvisitors import shorten_compound_type_annotation, ClassVisitor
 from py2puml.parsing.parseclassconstructor import parse_class_constructor
 from py2puml.parsing.moduleresolver import ModuleResolver
-# from py2puml.utils import investigate_domain_definition
 
 
 CONCRETE_TYPE_PATTERN = re_compile("^<(?:class|enum) '([\\.|\\w]+)'>$")
 
+
 def handle_inheritance_relation(
     class_type: Type,
     class_fqn: str,
@@ -28,6 +28,7 @@ def handle_inheritance_relation(
                 UmlRelation(base_type_fqn, class_fqn, RelType.INHERITANCE)
             )
 
+
 def inspect_static_attributes(
     class_type: Type,
     class_type_fqn: str,
@@ -59,7 +60,7 @@ def inspect_static_attributes(
         # utility which outputs the fully-qualified name of the attribute types
         module_resolver = ModuleResolver(import_module(class_type.__module__))
 
-        # builds the definitions of the class attrbutes and their relationships by iterating over the type annotations 
+        # builds the definitions of the class attrbutes and their relationships by iterating over the type annotations
         for attr_name, attr_class in type_annotations.items():
             attr_raw_type = str(attr_class)
             concrete_type_match = CONCRETE_TYPE_PATTERN.search(attr_raw_type)
@@ -89,18 +90,16 @@ def inspect_static_attributes(
 
     return definition_attrs
 
-def inspect_methods(
-    definition_methods, class_type,
-):
-    no_dunder = lambda method_name: not (method_name[0].startswith('__') or method_name[0].endswith('__'))
-    methods = filter(no_dunder, getmembers(class_type, callable))
-    for name, method in methods:
-        method_signature = signature(method)
-        uml_method = UmlMethod(
-            name=name,
-            signature=str(method_signature),
-        )
-        definition_methods.append(uml_method)
+
+def inspect_methods(definition_methods: List, class_type: Type, root_module_name: str):
+    """ This function parses a class using AST to identify methods. """
+    print(f'inspecting {class_type.__name__} from {class_type.__module__}')
+    class_source: str = getsource(class_type)
+    class_ast: AST = parse(class_source)
+    visitor = ClassVisitor(class_type, root_module_name)
+    visitor.visit(class_ast)
+    for method in visitor.uml_methods:
+        definition_methods.append(method)
 
 
 def inspect_class_type(
@@ -117,11 +116,12 @@ def inspect_class_type(
     instance_attributes, compositions = parse_class_constructor(class_type, class_type_fqn, root_module_name)
     attributes.extend(instance_attributes)
     domain_relations.extend(compositions.values())
-    
-    inspect_methods(domain_items_by_fqn[class_type_fqn].methods, class_type)
+
+    inspect_methods(domain_items_by_fqn[class_type_fqn].methods, class_type, root_module_name)
 
     handle_inheritance_relation(class_type, class_type_fqn, root_module_name, domain_relations)
 
+
 def inspect_dataclass_type(
     class_type: Type,
     class_type_fqn: str,

py2puml/parsing/astvisitors.py

@@ -1,36 +1,38 @@
 
-from typing import Dict, List, Tuple
-
+from typing import Dict, List, Tuple, Type
+from inspect import getsource, unwrap
 from ast import (
     NodeVisitor, arg, expr,
     FunctionDef, Assign, AnnAssign,
     Attribute, Name, Subscript, get_source_segment
 )
 from collections import namedtuple
+from textwrap import dedent
+from importlib import import_module
 
-from py2puml.domain.umlclass import UmlAttribute
+from py2puml.domain.umlclass import UmlAttribute, UmlMethod
 from py2puml.domain.umlrelation import UmlRelation, RelType
 from py2puml.parsing.compoundtypesplitter import CompoundTypeSplitter, SPLITTING_CHARACTERS
 from py2puml.parsing.moduleresolver import ModuleResolver, NamespacedType
 
-
 Variable = namedtuple('Variable', ['id', 'type_expr'])
 
+
 class SignatureVariablesCollector(NodeVisitor):
-    '''
-    Collects the variables and their type annotations from the signature of a constructor method
-    '''
-    def __init__(self, constructor_source: str, *args, **kwargs):
+    """
+    Collects the variables and their type annotations from the signature of a method
+    """
+    def __init__(self, skip_self=False, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.constructor_source = constructor_source
+        self.skip_self = skip_self
         self.class_self_id: str = None
         self.variables: List[Variable] = []
 
     def visit_arg(self, node: arg):
         variable = Variable(node.arg, node.annotation)
 
         # first constructor variable is the name for the 'self' reference
-        if self.class_self_id is None:
+        if self.class_self_id is None and not self.skip_self:
             self.class_self_id = variable.id
         # other arguments are constructor parameters
         else:
@@ -66,6 +68,57 @@ def visit_Subscript(self, node: Subscript):
         pass
 
 
+class ClassVisitor(NodeVisitor):
+
+    def __init__(self, class_type: Type, root_fqn: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.uml_methods: List[UmlMethod] = []
+
+    def visit_FunctionDef(self, node: FunctionDef):
+        method_visitor = MethodVisitor()
+        method_visitor.visit(node)
+        self.uml_methods.append(method_visitor.uml_method)
+
+
+class MethodVisitor(NodeVisitor):
+    """
+    Node visitor subclass used to walk the abstract syntax tree of a method class and identify method arguments.
+
+    If the method is the class constructor, instance attributes (and their type) are also identified by looking both at the constructor signature and constructor's body. When searching in the constructor's body, the visitor looks for relevant assignments (with and without type annotation).
+     """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.variables_namespace: List[Variable] = []
+        self.class_self_id: str
+        self.uml_method: UmlMethod
+
+    def generic_visit(self, node):
+        NodeVisitor.generic_visit(self, node)
+
+    def visit_FunctionDef(self, node: FunctionDef):
+        decorators = [decorator.id for decorator in node.decorator_list]
+        is_static = 'staticmethod' in decorators
+        is_class = 'classmethod' in decorators
+        variables_collector = SignatureVariablesCollector(skip_self=is_static)
+        variables_collector.visit(node)
+        self.variables_namespace = variables_collector.variables
+
+        if node.name == '__init__':
+            self.class_self_id: str = variables_collector.class_self_id
+            self.generic_visit(node)    #Only visit child nodes for constructor
+
+        self.uml_method = UmlMethod(name=node.name, is_static=is_static, is_class=is_class)
+        for argument in variables_collector.variables:
+            if argument.type_expr:
+                if hasattr(argument.type_expr, 'id'):
+                    self.uml_method.arguments[argument.id] = argument.type_expr.id
+                else:
+                    self.uml_method.arguments[argument.id] = f'SUBscript {argument.type_expr.value.id}'
+            else:
+                self.uml_method.arguments[argument.id] = None
+
+
 class ConstructorVisitor(NodeVisitor):
     '''
     Identifies the attributes (and infer their type) assigned to self in the body of a constructor method
@@ -105,7 +158,7 @@ def generic_visit(self, node):
     def visit_FunctionDef(self, node: FunctionDef):
         # retrieves constructor arguments ('self' reference and typed arguments)
         if node.name == '__init__':
-            variables_collector = SignatureVariablesCollector(self.constructor_source)
+            variables_collector = SignatureVariablesCollector()
             variables_collector.visit(node)
             self.class_self_id: str = variables_collector.class_self_id
             self.variables_namespace = variables_collector.variables
@@ -158,7 +211,6 @@ def visit_Assign(self, node: Assign):
             # other assignments were done in new variables that can shadow existing ones
             self.variables_namespace.extend(variables_collector.variables)
 
-
     def derive_type_annotation_details(self, annotation: expr) -> Tuple[str, List[str]]:
         '''
         From a type annotation, derives:
@@ -182,13 +234,12 @@ def derive_type_annotation_details(self, annotation: expr) -> Tuple[str, List[st
             return short_type, [full_namespaced_type]
         # compound type (List[...], Tuple[Dict[str, float], module.DomainType], etc.)
         elif isinstance(annotation, Subscript):
-            return shorten_compound_type_annotation(
-                get_source_segment(self.constructor_source, annotation),
-                self.module_resolver
-            )
-
+            source_segment = get_source_segment(self.constructor_source, annotation)
+            short_type, associated_types = shorten_compound_type_annotation(source_segment, self.module_resolver)
+            return short_type, associated_types
         return None, []
 
+
 def shorten_compound_type_annotation(type_annotation: str, module_resolver: ModuleResolver) -> Tuple[str, List[str]]:
     '''
     In the string representation of a compound type annotation, the elementary types can be prefixed by their packages or sub-packages.

tests/modules/withmethods/withmethods.py

@@ -1,11 +1,11 @@
+from __future__ import annotations
 from typing import List, Tuple
 from math import pi
 
 from tests.modules import withenum
 from tests import modules
 from tests.modules.withenum import TimeUnit
 
-
 class Coordinates:
     def __init__(self, x: float, y: float) -> None:
         self.x = x
@@ -30,3 +30,6 @@ def __init__(self, x: int, y: str) -> None:
         self.time_resolution: Tuple[str, withenum.TimeUnit] = 'minute', TimeUnit.MINUTE
         self.x = x
         self.y = y
+
+    def do_something(self, posarg_nohint, posarg_hint: str, posarg_default=3) -> int:
+        return 44

tests/puml_files/with_methods.puml

@@ -11,20 +11,20 @@ class tests.modules.withmethods.withmethods.Point {
   time_resolution: Tuple[str, TimeUnit]
   x: int
   y: str
-  {static} Point from_values(x: int, y: str)
+  {static} Point from_values(int x, str y)
   Tuple[float, str] get_coordinates(self)
-  __init__(self, x: int, y: str)
+  __init__(self, int x, str y)
 }
 class tests.modules.withmethods.withinheritedmethods.ThreeDimensionalPoint {
   z: float
-  __init__(self, x: int, y: str, z: float)
-  move(self, offset: int)
+  __init__(self, int x, str y, float z)
+  move(self, int offset)
   bool check_positive(self)
 }
 class tests.modules.withmethods.withmethods.Coordinates {
   x: float
   y: float
-  __init__(self, x: float, y: float)
+  __init__(self, float x, float y)
 }
 tests.modules.withmethods.withmethods.Point *-- tests.modules.withmethods.withmethods.Coordinates
 tests.modules.withmethods.withmethods.Point <|-- tests.modules.withmethods.withinheritedmethods.ThreeDimensionalPoint

tests/py2puml/inspection/test_inspectclass.py

@@ -172,14 +172,9 @@ def test_inspect_module_should_handle_compound_types_with_numbers_in_their_name(
 def test_inspect_module_should_find_methods(
     domain_items_by_fqn: Dict[str, UmlItem], domain_relations: List[UmlRelation]
 ):
-    '''
+    """
     Test that methods are detected including static methods
-
-    This test case assumes that the methods will be sorted by type as follow:
-        1a - instance methods (special methods aka "dunder")
-        1b - all other instance methods
-        2 - static methods
-    '''
+    """
 
     inspect_module(
         import_module('tests.modules.withmethods.withmethods'),
@@ -189,16 +184,16 @@ def test_inspect_module_should_find_methods(
 
     # Coordinates UmlClass
     coordinates_umlitem: UmlClass = domain_items_by_fqn['tests.modules.withmethods.withmethods.Coordinates']
-    assert len(coordinates_umlitem.methods) == 0
+    assert len(coordinates_umlitem.methods) == 1
 
     # Point UmlClass
     point_umlitem: UmlClass = domain_items_by_fqn['tests.modules.withmethods.withmethods.Point']
-    assert len(point_umlitem.methods) == 2
+    assert len(point_umlitem.methods) == 4
 
-    # FIXME dunder methods are filtered out for now
-    # assert point_umlitem.methods[0].name == '__init__'  # 1a - Instance method (special)
-    assert point_umlitem.methods[0].name == 'from_values'  # 2 - Static method
-    assert point_umlitem.methods[1].name == 'get_coordinates'  # 1b - Instance method (regular)
+    assert point_umlitem.methods[0].name == 'from_values'
+    assert point_umlitem.methods[1].name == 'get_coordinates'
+    assert point_umlitem.methods[2].name == '__init__'
+    assert point_umlitem.methods[3].name == 'do_something'
     # FIXME: use 'assert_method' once UmlMethod restructured
 
 
@@ -219,10 +214,9 @@ def test_inspect_module_inherited_methods(
     coordinates_3d_umlitem: UmlClass = domain_items_by_fqn['tests.modules.withmethods.withinheritedmethods.ThreeDimensionalPoint']
 
     # FIXME inherited methods should not be mentionned
-    assert len(coordinates_3d_umlitem.methods) == 4
+    assert len(coordinates_3d_umlitem.methods) == 3
 
-    assert coordinates_3d_umlitem.methods[0].name == 'check_positive'
-    assert coordinates_3d_umlitem.methods[1].name == 'from_values' # inherited method
-    assert coordinates_3d_umlitem.methods[2].name == 'get_coordinates' # inherited method
-    assert coordinates_3d_umlitem.methods[3].name == 'move'
+    assert coordinates_3d_umlitem.methods[2].name == 'check_positive'
+    assert coordinates_3d_umlitem.methods[0].name == '__init__'
+    assert coordinates_3d_umlitem.methods[1].name == 'move'
     # FIXME: use 'assert_method' once UmlMethod restructured

tests/py2puml/parsing/test_astvisitors.py

@@ -12,6 +12,7 @@
 
 from tests.asserts.variable import assert_Variable
 from tests.py2puml.parsing.mockedinstance import MockedInstance
+from tests.modules.withmethods import withmethods
 
 
 class ParseMyConstructorArguments:
@@ -31,7 +32,7 @@ def test_SignatureVariablesCollector_collect_arguments():
     constructor_source: str = dedent(getsource(ParseMyConstructorArguments.__init__.__code__))
     constructor_ast: AST = parse(constructor_source)
 
-    collector = SignatureVariablesCollector(constructor_source)
+    collector = SignatureVariablesCollector()
     collector.visit(constructor_ast)
 
     assert collector.class_self_id == 'me'