Allow users to choose the name strategy of fields

protobuf_uml_diagram.py

@@ -38,7 +38,8 @@
 class PathPath(click.Path):
     """A Click path argument that returns a pathlib Path, not a string"""
 
-    def convert(self, value: Union[str, PathLike], param: Optional[click.Parameter], ctx: Union[click.Context, None]) -> Path:
+    def convert(self, value: Union[str, PathLike], param: Optional[click.Parameter],
+                ctx: Union[click.Context, None]) -> Path:
         """Convert a text parameter into a ``Path`` object.
         :param value: parameter value
         :type value: str
@@ -61,54 +62,71 @@ def convert(self, value: Union[str, PathLike], param: Optional[click.Parameter],
     for text, number in FieldDescriptorProto.Type.items()
 }
 
-
 LABELS_BY_NUMBER = {
     number: text.lower().replace("label_", "")
     for text, number in FieldDescriptorProto.Label.items()
 }
 
 
-def _process_module(proto_module: ModuleType) -> Tuple[List[str], List[str]]:
+def _process_module(proto_module: ModuleType, full_names=True) -> Tuple[List[str], List[str]]:
     """"
     :return: list of descriptors
     :rtype: List[Descriptor]
+    :param full_names: whether the output must include the full name of classes
+    :type full_names: bool
     """
     classes: List[str] = []
     relationships: List[str] = []
     for type_name, type_descriptor in proto_module.DESCRIPTOR.message_types_by_name.items():
-        _process_descriptor(type_descriptor, classes, relationships)
+        _process_descriptor(type_descriptor, classes, relationships, full_names=full_names)
     return classes, relationships
 
 
-def _process_descriptor(descriptor: Descriptor, classes: List[str],
-                        relationships: List[str]) -> None:
+def _get_field_name(descriptor: Descriptor, full_names=True) -> str:
+    if full_names:
+        return descriptor.full_name
+    return descriptor.name
+
+
+
+def _process_descriptor(
+        descriptor: Descriptor, classes: List[str],
+        relationships: List[str],
+        full_names=True) -> None:
     """
     :param descriptor: a Protobuf descriptor
     :type descriptor: Descriptor
     :param classes: list of classes
     :type classes: list
+    :param full_names: whether the output must include the full name of classes
+    :type full_names: bool
     """
+    # Here users are able to choose between ClassName.type_name (full name included) or just type_name
+
+
     type_template_text = StringIO()
-    this_node = descriptor.full_name
+    this_node = _get_field_name(descriptor, full_names=full_names)
     type_template_text.write(
         f"""    \"{this_node}\"[label = "{{{this_node}|""")
     fields = []
     for _field in descriptor.fields:
         if _field.type == FieldDescriptor.TYPE_MESSAGE:
-            that_node = _field.message_type.full_name
+            that_node = _get_field_name(_field.message_type, full_names=full_names)
 
             # is it a repeated field?
             label = LABELS_BY_NUMBER[_field.label]
             if label == 'repeated':
-                relationships.append(f"    \"{that_node}\"->\"{this_node}\" [dir=backward;arrowhead=odiamond,arrowtail=normal;headlabel=\"1\";taillabel=\"0..*\"]")
+                relationships.append(
+                    f"    \"{that_node}\"->\"{this_node}\" [dir=backward;arrowhead=odiamond,arrowtail=normal;headlabel=\"1\";taillabel=\"0..*\"]")
             else:
-                relationships.append(f"    \"{this_node}\"->\"{that_node}\" [arrowhead=none;headlabel=\"1\";taillabel=\"1\"]")
+                relationships.append(
+                    f"    \"{this_node}\"->\"{that_node}\" [arrowhead=none;headlabel=\"1\";taillabel=\"1\"]")
 
             field_type = that_node  # so we replace the 'message' token by the actual name
         else:
             field_type = TYPES_BY_NUMBER[_field.type]
 
-        fields.append(f"+ {_field.full_name}:{field_type}")
+        fields.append(f"+ {_get_field_name(_field, full_names=full_names)}:{field_type}")
 
     # add fields
     type_template_text.write("\\n".join(fields))
@@ -119,19 +137,21 @@ def _process_descriptor(descriptor: Descriptor, classes: List[str],
 
     # nested types
     for nested_descriptor in descriptor.nested_types:
-        _process_descriptor(nested_descriptor, classes, relationships)
+        _process_descriptor(nested_descriptor, classes, relationships, full_names=full_names)
     # TODO: what about extension, enum, ...?
 
 
-def _get_uml_template(proto_module: ModuleType) -> str:
+def _get_uml_template(proto_module: ModuleType, full_names=True) -> str:
     """
     Return the graphviz dot template for a UML class diagram.
     :param proto_module: protobuf module
     :type proto_module: ModuleType
+    :param full_names: whether the output must include the full name of classes
+    :type full_names: bool
     :return: UML template
     :rtype: str
     """
-    classes, relationships = _process_module(proto_module)
+    classes, relationships = _process_module(proto_module, full_names=full_names)
     uml_template = Template("""
 digraph "Protobuf UML class diagram" {
     fontname="Bitstream Vera Sans"
@@ -176,6 +196,7 @@ class Diagram:
     _proto_module: Union[ModuleType, None] = None
     _rendered_filename: Union[str, None] = None
     _file_format = "png"
+    _full_names = True
 
     def from_file(self, proto_file: str):
         if not proto_file:
@@ -203,6 +224,10 @@ def with_format(self, file_format: str):
         self._file_format = file_format
         return self
 
+    def with_full_names(self, full_names: bool):
+        self._full_names = full_names
+        return self
+
     def build(self):
         if not self._proto_module:
             raise ValueError("No Protobuf Python module!")
@@ -211,7 +236,7 @@ def build(self):
         if not self._file_format:
             raise ValueError("No file format!")
 
-        uml_template = _get_uml_template(self._proto_module)
+        uml_template = _get_uml_template(self._proto_module, full_names=self._full_names)
 
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug("UML template:")
@@ -231,10 +256,12 @@ def build(self):
               help='Compiled Python proto module (e.g. some.package.ws_compiled_pb2).')
 @click.option('--output', type=PathPath(file_okay=False), required=True,
               help='Output directory.')
-def main(proto: str, output: Path) -> None:
+@click.option('--full_names', type=bool, required=False, default=True, help='Use full names (Class.type) or not (type) in diagram.')
+def main(proto: str, output: Path, full_names: bool) -> None:
     Diagram() \
         .from_file(proto) \
         .to_file(output) \
+        .with_full_names(full_names) \
         .build()