Add mypy to CI and fix issues in the code

.github/workflows/main.yml

@@ -45,6 +45,9 @@ jobs:
     - name: Linter
       run: |
         pycodestyle --filename=protobuf-uml-diagram --exclude=.git,__pycache__,.tox,venv,protobuf_uml_diagram.egg-info,.pytest_cache --max-line-length=120
+    - name: MyPy
+      run: |
+        mypy protobuf_uml_diagram.py
     - name: Unit tests with coverage
       run: |
         coverage run -p setup.py test

CHANGES.md

@@ -12,6 +12,7 @@
 - Update pytest-env requirement from <0.9,>=0.6 to >=0.6,<1.1 #72
 - Update pytest-env requirement from <1.1,>=0.6 to >=0.6,<1.2 #73
 - Update protobuf requirement from <4.25,>=3.13 to >=3.13,<4.26 #74
+- Add mypy #76
 
 ## 0.12 (11/03/2023)
 

protobuf_uml_diagram.py

@@ -19,38 +19,38 @@
 import logging
 from importlib import import_module
 from io import StringIO
+from os import PathLike
 from pathlib import Path
 from string import Template
 from types import ModuleType
-from typing import List, Tuple, Union
+from typing import cast, List, Optional, Tuple, Union
 
 import click
 from google.protobuf.descriptor import Descriptor, FieldDescriptor
 from google.protobuf.descriptor_pb2 import FieldDescriptorProto
-from graphviz import Source
+from graphviz import Source  # type: ignore  # TODO: https://github.com/xflr6/graphviz/issues/203
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
-Text = Union[str, bytes]
-
 
 # https://github.com/pallets/click/issues/405#issuecomment-470812067
 class PathPath(click.Path):
     """A Click path argument that returns a pathlib Path, not a string"""
 
-    def convert(self, value: Text, param: Text, ctx) -> 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: Text
+        :type value: str
         :param param: parameter name
-        :type param: Text
+        :type param: str
         :param ctx: context
-        :type ctx: object
+        :type ctx: click.Context
         :return: a ``Path`` object
         :rtype: Path
         """
-        return Path(super().convert(value, param, ctx))
+        p = super().convert(value, param, ctx)
+        return Path(cast(Path, p))
 
 
 # -- UML diagram
@@ -73,15 +73,15 @@ def _process_module(proto_module: ModuleType) -> Tuple[List[str], List[str]]:
     :return: list of descriptors
     :rtype: List[Descriptor]
     """
-    classes = []
-    relationships = []
+    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)
     return classes, relationships
 
 
-def _process_descriptor(descriptor: Descriptor, classes: list,
-                        relationships: list) -> None:
+def _process_descriptor(descriptor: Descriptor, classes: List[str],
+                        relationships: List[str]) -> None:
     """
     :param descriptor: a Protobuf descriptor
     :type descriptor: Descriptor
@@ -173,8 +173,8 @@ def _module(proto: str) -> ModuleType:
 class Diagram:
     """A diagram builder."""
 
-    _proto_module: ModuleType = None
-    _rendered_filename: str = None
+    _proto_module: Union[ModuleType, None] = None
+    _rendered_filename: Union[str, None] = None
     _file_format = "png"
 
     def from_file(self, proto_file: str):
@@ -191,6 +191,8 @@ def from_file(self, proto_file: str):
     def to_file(self, output: Path):
         if not output:
             raise ValueError("Missing output location!")
+        if not self._proto_module or not self._proto_module.__file__:
+            raise ValueError("Missing protobuf module!")
         uml_file = Path(self._proto_module.__file__).stem
         self._rendered_filename = str(output.joinpath(uml_file))
         return self

setup.py

@@ -15,7 +15,7 @@
 import codecs
 from os.path import join, dirname, abspath
 
-from setuptools import setup
+from setuptools import setup  # type: ignore
 
 here = abspath(dirname(__file__))
 
@@ -35,6 +35,11 @@ def read(*parts):
     'pytest-runner>=4.1,<6.1'
 ]
 
+mypy_requires = [
+    'mypy==1.*',
+    'types-protobuf==4.24.*'
+]
+
 tests_require = [
     'codecov==2.1.*',
     'coverage>=5.3,<7.4',
@@ -46,7 +51,7 @@ def read(*parts):
 
 extras_require = {
     'tests': tests_require,
-    'all': install_requires + tests_require
+    'all': install_requires + tests_require + mypy_requires
 }
 
 setup(

tests.py

@@ -37,12 +37,12 @@ def test_from_file_raises(self):
 
     def test_to_file_raises(self):
         with pytest.raises(ValueError) as e:
-            Diagram().to_file(None)
+            Diagram().to_file(None)  # type: ignore
         assert 'Missing output location' in str(e.value)
 
     def test_with_format_raises(self):
         with pytest.raises(ValueError) as e:
-            Diagram().with_format(None)
+            Diagram().with_format(None)  # type: ignore
         assert 'Missing file' in str(e.value)
 
     def test_build_raises(self):
@@ -103,3 +103,27 @@ def test_contains_dot_proto_in_middle_of_the_name(self):
                 .build()
             assert os.path.getsize(tf) > 0
 
+
+    def test_to_file_with_missing_protobuf(self):
+        """A test where the protobuf module is missing."""
+        with TemporaryDirectory() as tmpdir:
+            tf = os.path.join(tmpdir, 'diagram.png')
+            d = Diagram().from_file('test_data.issue_27.proto.configs_data_pb2')
+            d._proto_module = None
+            with pytest.raises(ValueError) as cm:
+                d.to_file(Path(tf))
+            assert str(cm.value) == 'Missing protobuf module!'
+
+
+    def test_to_file_with_protobuf_missing_file(self):
+        """A test where the protobuf module is present but invalid (no file)."""
+        with TemporaryDirectory() as tmpdir:
+            tf = os.path.join(tmpdir, 'diagram.png')
+            d = Diagram().from_file('test_data.issue_27.proto.configs_data_pb2')
+            d._proto_module = {}
+            with pytest.raises(ValueError) as cm:
+                d.to_file(Path(tf))
+            assert str(cm.value) == 'Missing protobuf module!'
+
+
+