Add support for compound type

Rename named tuple Variable to Argument in astvisitors.py
Rename ReturnTypeVisitor to TypeVisitor which now support compound
type

Author: grayfox88

py2puml/parsing/astvisitors.py

@@ -1,4 +1,4 @@
-
+import ast
 from typing import Dict, List, Tuple, Type
 from ast import (
     NodeVisitor, arg, expr,
@@ -12,50 +12,56 @@
 from py2puml.parsing.compoundtypesplitter import CompoundTypeSplitter, SPLITTING_CHARACTERS
 from py2puml.parsing.moduleresolver import ModuleResolver
 
-Variable = namedtuple('Variable', ['id', 'type_expr'])
+Argument = namedtuple('Argument', ['id', 'type_expr'])
 
 
-class SignatureVariablesCollector(NodeVisitor):
+class SignatureArgumentsCollector(NodeVisitor):
     """
-    Collects the variables and their type annotations from the signature of a method
+    Collects the arguments name and type annotations from the signature of a method
     """
     def __init__(self, skip_self=False, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.skip_self = skip_self
         self.class_self_id: str = None
-        self.variables: List[Variable] = []
+        self.arguments: List[Argument] = []
+        self.datatypes = {}
 
     def visit_arg(self, node: arg):
-        variable = Variable(node.arg, node.annotation)
-
-        # first constructor variable is the name for the 'self' reference
+        argument = Argument(node.arg, node.annotation)
+        if node.annotation:
+            type_visitor = TypeVisitor()
+            datatype = type_visitor.visit(node.annotation)
+        else:
+            datatype = None
+        self.datatypes[node.arg] = datatype
+        # first constructor argument is the name for the 'self' reference
         if self.class_self_id is None and not self.skip_self:
-            self.class_self_id = variable.id
+            self.class_self_id = argument.id
         # other arguments are constructor parameters
-        self.variables.append(variable)
+        self.arguments.append(argument)
 
 
 class AssignedVariablesCollector(NodeVisitor):
     '''Parses the target of an assignment statement to detect whether the value is assigned to a variable or an instance attribute'''
     def __init__(self, class_self_id: str, annotation: expr):
         self.class_self_id: str = class_self_id
         self.annotation: expr = annotation
-        self.variables: List[Variable] = []
-        self.self_attributes: List[Variable] = []
+        self.variables: List[Argument] = []
+        self.self_attributes: List[Argument] = []
 
     def visit_Name(self, node: Name):
         '''
         Detects declarations of new variables
         '''
         if node.id != self.class_self_id:
-            self.variables.append(Variable(node.id, self.annotation))
+            self.variables.append(Argument(node.id, self.annotation))
 
     def visit_Attribute(self, node: Attribute):
         '''
         Detects declarations of new attributes on 'self'
         '''
         if isinstance(node.value, Name) and node.value.id == self.class_self_id:
-            self.self_attributes.append(Variable(node.attr, self.annotation))
+            self.self_attributes.append(Argument(node.attr, self.annotation))
 
     def visit_Subscript(self, node: Subscript):
         '''
@@ -76,8 +82,8 @@ def visit_FunctionDef(self, node: FunctionDef):
         self.uml_methods.append(method_visitor.uml_method)
 
 
-class ReturnTypeVisitor(NodeVisitor):
-
+class TypeVisitor(NodeVisitor):
+    """ Returns a string representation of a data type. Supports nested compound data types """
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -87,8 +93,24 @@ def visit_Name(self, node):
     def visit_Constant(self, node):
         return node.value
 
-    def visit_Subscript(self, node):
-        return node.value.id
+    def visit_Subscript(self, node: Subscript):
+        """ Visit node of type ast.Subscript and returns a string representation of the compound datatype. Iterate
+        over elements contained in slice attribute by calling recursively visit() method of new instances of
+        TypeVisitor. This allows to resolve nested compound datatype. """
+
+        datatypes = []
+
+        if hasattr(node.slice.value, 'elts'):
+            for child_node in node.slice.value.elts:
+                child_visitor = TypeVisitor()
+                datatypes.append(child_visitor.visit(child_node))
+        else:
+            child_visitor = TypeVisitor()
+            datatypes.append(child_visitor.visit(node.slice.value))
+
+        joined_datatypes = ', '.join(datatypes)
+
+        return f'{node.value.id}[{joined_datatypes}]'
 
 
 class MethodVisitor(NodeVisitor):
@@ -100,32 +122,33 @@ class MethodVisitor(NodeVisitor):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.variables_namespace: List[Variable] = []
+        self.variables_namespace: List[Argument] = []
         self.uml_method: UmlMethod
 
     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
+        arguments_collector = SignatureArgumentsCollector(skip_self=is_static)
+        arguments_collector.visit(node)
+        self.variables_namespace = arguments_collector.arguments
 
         self.uml_method = UmlMethod(name=node.name, is_static=is_static, is_class=is_class)
 
-        for argument in variables_collector.variables:
-            if argument.id == variables_collector.class_self_id:
+        for argument in arguments_collector.arguments:
+            if argument.id == arguments_collector.class_self_id:
                 self.uml_method.arguments[argument.id] = None
             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}' #FIXME
+
+                    self.uml_method.arguments[argument.id] = arguments_collector.datatypes[argument.id]
             else:
                 self.uml_method.arguments[argument.id] = None
 
         if node.returns is not None:
-            return_visitor = ReturnTypeVisitor()
+            return_visitor = TypeVisitor()
             self.uml_method.return_type = return_visitor.visit(node.returns)
 
 
@@ -140,7 +163,7 @@ def __init__(self, constructor_source: str, class_name: str, root_fqn: str, modu
         self.root_fqn = root_fqn
         self.module_resolver = module_resolver
         self.class_self_id: str
-        self.variables_namespace: List[Variable] = []
+        self.variables_namespace: List[Argument] = []
         self.uml_attributes: List[UmlAttribute] = []
         self.uml_relations_by_target_fqn: Dict[str, UmlRelation] = {}
 
@@ -153,7 +176,7 @@ def extend_relations(self, target_fqns: List[str]):
             )
         })
 
-    def get_from_namespace(self, variable_id: str) -> Variable:
+    def get_from_namespace(self, variable_id: str) -> Argument:
         return next((
             variable
             # variables namespace is iterated antichronologically
@@ -168,10 +191,10 @@ 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()
+            variables_collector = SignatureArgumentsCollector()
             variables_collector.visit(node)
             self.class_self_id: str = variables_collector.class_self_id
-            self.variables_namespace = variables_collector.variables
+            self.variables_namespace = variables_collector.arguments
 
         self.generic_visit(node)
 

tests/asserts/variable.py

@@ -1,9 +1,9 @@
 
 from ast import get_source_segment
 
-from py2puml.parsing.astvisitors import Variable
+from py2puml.parsing.astvisitors import Argument
 
-def assert_Variable(variable: Variable, id: str, type_str: str, source_code: str):
+def assert_Variable(variable: Argument, id: str, type_str: str, source_code: str):
     assert variable.id == id
     if type_str is None:
         assert variable.type_expr == None, 'no type annotation'

tests/modules/withmethods/withmethods.py

@@ -6,6 +6,7 @@
 from tests import modules
 from tests.modules.withenum import TimeUnit
 
+
 class Coordinates:
     def __init__(self, x: float, y: float) -> None:
         self.x = x
@@ -23,7 +24,7 @@ def from_values(x: int, y: str) -> Point:
     def get_coordinates(self) -> Tuple[float, str]:
         return self.x, self.y
 
-    def __init__(self, x: int, y: str) -> None:
+    def __init__(self, x: int, y: Tuple[bool]) -> None:
         self.coordinates: Coordinates = Coordinates(x, float(y))
         self.day_unit: withenum.TimeUnit = withenum.TimeUnit.DAYS
         self.hour_unit: modules.withenum.TimeUnit = modules.withenum.TimeUnit.HOURS

tests/puml_files/with_methods.puml

@@ -10,10 +10,10 @@ class tests.modules.withmethods.withmethods.Point {
   hour_unit: TimeUnit
   time_resolution: Tuple[str, TimeUnit]
   x: int
-  y: str
+  y: Tuple[bool]
   {static} Point from_values(int x, str y)
   Tuple[float, str] get_coordinates(self)
-  __init__(self, int x, str y)
+  __init__(self, int x, Tuple[bool] y)
   int do_something(self, posarg_nohint, str posarg_hint, posarg_default)
 }
 class tests.modules.withmethods.withinheritedmethods.ThreeDimensionalPoint {

tests/py2puml/parsing/test_astvisitors.py

@@ -1,4 +1,4 @@
-
+import unittest
 from typing import Dict, Tuple, List
 
 from ast import parse, AST, get_source_segment
@@ -7,7 +7,7 @@
 
 from pytest import mark
 
-from py2puml.parsing.astvisitors import AssignedVariablesCollector, SignatureVariablesCollector, Variable, shorten_compound_type_annotation
+from py2puml.parsing.astvisitors import AssignedVariablesCollector, TypeVisitor, SignatureArgumentsCollector, Argument, shorten_compound_type_annotation
 from py2puml.parsing.moduleresolver import ModuleResolver
 
 from tests.asserts.variable import assert_Variable
@@ -32,18 +32,18 @@ def test_SignatureVariablesCollector_collect_arguments():
     constructor_source: str = dedent(getsource(ParseMyConstructorArguments.__init__.__code__))
     constructor_ast: AST = parse(constructor_source)
 
-    collector = SignatureVariablesCollector()
+    collector = SignatureArgumentsCollector()
     collector.visit(constructor_ast)
 
     assert collector.class_self_id == 'me'
-    assert len(collector.variables) == 7, 'all the arguments must be detected'
-    assert_Variable(collector.variables[0], 'me', None, constructor_source)
-    assert_Variable(collector.variables[1], 'an_int', 'int', constructor_source)
-    assert_Variable(collector.variables[2], 'an_untyped', None, constructor_source)
-    assert_Variable(collector.variables[3], 'a_compound_type', 'Tuple[float, Dict[str, List[bool]]]', constructor_source)
-    assert_Variable(collector.variables[4], 'a_default_string', 'str', constructor_source)
-    assert_Variable(collector.variables[5], 'args', None, constructor_source)
-    assert_Variable(collector.variables[6], 'kwargs', None, constructor_source)
+    assert len(collector.arguments) == 7, 'all the arguments must be detected'
+    assert_Variable(collector.arguments[0], 'me', None, constructor_source)
+    assert_Variable(collector.arguments[1], 'an_int', 'int', constructor_source)
+    assert_Variable(collector.arguments[2], 'an_untyped', None, constructor_source)
+    assert_Variable(collector.arguments[3], 'a_compound_type', 'Tuple[float, Dict[str, List[bool]]]', constructor_source)
+    assert_Variable(collector.arguments[4], 'a_default_string', 'str', constructor_source)
+    assert_Variable(collector.arguments[5], 'args', None, constructor_source)
+    assert_Variable(collector.arguments[6], 'kwargs', None, constructor_source)
 
 @mark.parametrize(
     'class_self_id,assignment_code,annotation_as_str,self_attributes,variables', [
@@ -190,3 +190,45 @@ def test_shorten_compound_type_annotation(full_annotation: str, short_annotation
     shortened_annotation, full_namespaced_definitions = shorten_compound_type_annotation(full_annotation, module_resolver)
     assert shortened_annotation == short_annotation
     assert full_namespaced_definitions == namespaced_definitions
+
+
+class TestTypeVisitor(unittest.TestCase):
+
+    def test_return_type_int(self):
+        source_code = 'def dummy_function() -> int:\n     pass'
+        ast = parse(source_code)
+        node = ast.body[0].returns
+        visitor = TypeVisitor()
+        actual_rtype = visitor.visit(node)
+        expected_rtype = 'int'
+        self.assertEqual(expected_rtype, actual_rtype)
+
+    def test_return_type_compound(self):
+        """ Test non-nested compound datatype"""
+        source_code = 'def dummy_function() -> Tuple[float, str]:\n     pass'
+        ast = parse(source_code)
+        node = ast.body[0].returns
+        visitor = TypeVisitor()
+        actual_rtype = visitor.visit(node)
+        expected_rtype = 'Tuple[float, str]'
+        self.assertEqual(expected_rtype, actual_rtype)
+
+    def test_return_type_compound_nested(self):
+        """ Test nested compound datatype"""
+        source_code = 'def dummy_function() -> Tuple[float, Dict[str, List[bool]]]:\n     pass'
+        ast = parse(source_code)
+        node = ast.body[0].returns
+        visitor = TypeVisitor()
+        actual_rtype = visitor.visit(node)
+        expected_rtype = 'Tuple[float, Dict[str, List[bool]]]'
+        self.assertEqual(expected_rtype, actual_rtype)
+
+    def test_return_type_user_defined(self):
+        """ Test user-defined class datatype"""
+        source_code = 'def dummy_function() -> Point:\n     pass'
+        ast = parse(source_code)
+        node = ast.body[0].returns
+        visitor = TypeVisitor()
+        actual_rtype = visitor.visit(node)
+        expected_rtype = 'Point'
+        self.assertEqual(expected_rtype, actual_rtype)