Implement return_type

Implement return_type for methods based on type hinting
(not working with compound type). Correct test cases.

Fix bug where self was not included in method signature.

Remove useless import statements.

Refactor puml generation for class methods using

UmlMethod.represent_as_puml

Author: grayfox88

py2puml/domain/umlclass.py

@@ -10,6 +10,7 @@ class UmlAttribute:
     type: str
     static: bool
 
+
 @dataclass
 class UmlMethod:
     name: str
@@ -18,10 +19,19 @@ class UmlMethod:
     is_class: bool = False
     return_type: str = None
 
+    def represent_as_puml(self):
+        items = []
+        if self.is_static:
+            items.append('{static}')
+        if self.return_type:
+            items.append(self.return_type)
+        items.append(f'{self.name}({self.signature})')
+        return ' '.join(items)
+
     @property
     def signature(self):
         if self.arguments:
-            return ', '.join([f'{arg_type} {arg_name}' for arg_name, arg_type in self.arguments.items()])
+            return ', '.join([f'{arg_type} {arg_name}' if arg_type else f'{arg_name}' for arg_name, arg_type in self.arguments.items()])
         return ''
 
 

py2puml/export/puml.py

@@ -12,7 +12,6 @@
 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 = '  {staticity} {name}({signature})\n'
 PUML_ITEM_END = '}\n'
 PUML_RELATION_TPL = '{source_fqn} {rel_type}-- {target_fqn}\n'
 
@@ -40,7 +39,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, staticity=FEATURE_STATIC if uml_method.is_static else FEATURE_INSTANCE)
+                yield f'  {uml_method.represent_as_puml()}\n'
             yield PUML_ITEM_END
         else:
             raise TypeError(f'cannot process uml_item of type {uml_item.__class__}')

py2puml/parsing/astvisitors.py

@@ -1,19 +1,16 @@
 
 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, UmlMethod
 from py2puml.domain.umlrelation import UmlRelation, RelType
 from py2puml.parsing.compoundtypesplitter import CompoundTypeSplitter, SPLITTING_CHARACTERS
-from py2puml.parsing.moduleresolver import ModuleResolver, NamespacedType
+from py2puml.parsing.moduleresolver import ModuleResolver
 
 Variable = namedtuple('Variable', ['id', 'type_expr'])
 
@@ -35,8 +32,7 @@ def visit_arg(self, node: arg):
         if self.class_self_id is None and not self.skip_self:
             self.class_self_id = variable.id
         # other arguments are constructor parameters
-        else:
-            self.variables.append(variable)
+        self.variables.append(variable)
 
 
 class AssignedVariablesCollector(NodeVisitor):
@@ -80,6 +76,21 @@ def visit_FunctionDef(self, node: FunctionDef):
         self.uml_methods.append(method_visitor.uml_method)
 
 
+class ReturnTypeVisitor(NodeVisitor):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def visit_Name(self, node):
+        return node.id
+
+    def visit_Constant(self, node):
+        return node.value
+
+    def visit_Subscript(self, node):
+        return node.value.id
+
+
 class MethodVisitor(NodeVisitor):
     """
     Node visitor subclass used to walk the abstract syntax tree of a method class and identify method arguments.
@@ -90,12 +101,8 @@ class MethodVisitor(NodeVisitor):
     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
@@ -104,20 +111,23 @@ def visit_FunctionDef(self, node: FunctionDef):
         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.id == variables_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}'
+                    self.uml_method.arguments[argument.id] = f'Subscript {argument.type_expr.value.id}' #FIXME
             else:
                 self.uml_method.arguments[argument.id] = None
 
+        if node.returns is not None:
+            return_visitor = ReturnTypeVisitor()
+            self.uml_method.return_type = return_visitor.visit(node.returns)
+
 
 class ConstructorVisitor(NodeVisitor):
     '''

py2puml/py2puml.domain.puml

@@ -27,7 +27,10 @@ class py2puml.domain.umlitem.UmlItem {
 }
 class py2puml.domain.umlclass.UmlMethod {
   name: str
-  signature: str
+  arguments: Dict
+  is_static: bool
+  is_class: bool
+  return_type: str
 }
 class py2puml.domain.umlenum.Member {
   name: str

tests/modules/withmethods/withmethods.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import List, Tuple
+from typing import Tuple
 from math import pi
 
 from tests.modules import withenum
@@ -17,7 +17,7 @@ class Point:
     origin = Coordinates(0, 0)
 
     @staticmethod
-    def from_values(x: int, y: str) -> 'Point':
+    def from_values(x: int, y: str) -> Point:
         return Point(x, y)
 
     def get_coordinates(self) -> Tuple[float, str]:

tests/puml_files/with_methods.puml

@@ -14,6 +14,7 @@ class tests.modules.withmethods.withmethods.Point {
   {static} Point from_values(int x, str y)
   Tuple[float, str] get_coordinates(self)
   __init__(self, int x, str y)
+  int do_something(self, posarg_nohint, str posarg_hint, posarg_default)
 }
 class tests.modules.withmethods.withinheritedmethods.ThreeDimensionalPoint {
   z: float

tests/py2puml/parsing/test_astvisitors.py

@@ -36,13 +36,14 @@ def test_SignatureVariablesCollector_collect_arguments():
     collector.visit(constructor_ast)
 
     assert collector.class_self_id == 'me'
-    assert len(collector.variables) == 6, 'all the arguments must be detected'
-    assert_Variable(collector.variables[0], 'an_int', 'int', constructor_source)
-    assert_Variable(collector.variables[1], 'an_untyped', None, constructor_source)
-    assert_Variable(collector.variables[2], 'a_compound_type', 'Tuple[float, Dict[str, List[bool]]]', constructor_source)
-    assert_Variable(collector.variables[3], 'a_default_string', 'str', constructor_source)
-    assert_Variable(collector.variables[4], 'args', None, constructor_source)
-    assert_Variable(collector.variables[5], 'kwargs', None, constructor_source)
+    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)
 
 @mark.parametrize(
     'class_self_id,assignment_code,annotation_as_str,self_attributes,variables', [