Merge pull request #5 from tannewt/improve_types

 Improve union type handling and support :data:

.pylint

@@ -131,6 +131,7 @@ disable=print-statement,
         missing-docstring,
         invalid-name,
         global-statement,
+        too-many-branches,
         too-many-locals,
         too-many-instance-attributes,
         too-many-public-methods,

rst2pyi/core.py

@@ -16,7 +16,7 @@
 log = logging.getLogger(__name__)
 directive_re = re.compile(r"\s*..\s+(\w+)::\s+([^\n]+)")
 callable_re = re.compile(r"\s*(\w+)(?:\(([^\)]*)\))?")
-info_re = re.compile(r"\s*:(\w+)\s+[!~]?(.+)\s+(\w+):")
+info_re = re.compile(r"\s*:([a-z]+)\s+[!~]?([^:]+)\s+([^:]+):")
 
 
 def setup_logger(debug: bool = False):
@@ -103,6 +103,39 @@ def render(self, line: Line, kind: str = None, **kwargs: str) -> str:
             **kwargs,
         )
 
+    @staticmethod
+    def _split_types(t: str) -> Set[str]:
+        """Splits the types of all params into their individual pieces for import."""
+        out = set()
+        for typelist in t.split("["):
+            typelist = typelist.strip("]")
+            for subtype in typelist.split(","):
+                out.add(subtype.strip())
+        return out
+
+    @staticmethod
+    def _convert_type(t: str) -> str:
+        """
+        Legacy rST can use " or " to express type unions so convert it to PEP484 syntax.
+
+        Documented here:
+        http://www.sphinx-doc.org/en/master/usage/restructuredtext/domains.html#info-field-lists
+        """
+        if " or " in t:
+            types = t.split(" or ")
+            if "None" in types:
+                types.remove("None")
+                if len(types) > 1:
+                    return f"Optional[Union[{', '.join(types)}]]"
+                elif len(types) == 1:
+                    return f"Optional[{types[0]}]"
+                else:
+                    log.error("converting type string with no other types: %s", t)
+                    return t
+            else:
+                return f"Union[{', '.join(types)}]"
+        return t
+
     def gen_stub(self, dest: Path, lines: Lines) -> None:
         log.debug("generating stub %s", dest)
         config = self.config
@@ -136,6 +169,7 @@ def gen_stub(self, dest: Path, lines: Lines) -> None:
 
                 name, param_str = match.groups()
                 if param_str is None:
+                    log.warning("Missing param string %s %s:%d", call, path, lineno)
                     content.append(
                         self.render(line, name=name, args="", ret_type="Any")
                     )
@@ -152,12 +186,26 @@ def gen_stub(self, dest: Path, lines: Lines) -> None:
                 while idx + 1 < count and lines[idx + 1].kind == "param":
                     idx += 1
                     p_type, p_name = lines[idx].extra
+
+                    p_type = Converter._convert_type(p_type)
+                    matched = False
                     for pidx, (n, _, v) in enumerate(params):
                         if n == p_name:
                             params[pidx][1] = p_type
+                            matched = True
                             break
+                    if not matched:
+                        param = lines[idx]
+                        log.warning(
+                            "%s:%d: Param missing from function call: %s",
+                            param.source,
+                            param.lineno,
+                            param.extra[1],
+                        )
+
+                for _, t, _ in params:
+                    type_names.update(Converter._split_types(t))
 
-                type_names.update(t for _, t, _ in params)
                 args = ", ".join(
                     (
                         self.render(
@@ -179,7 +227,7 @@ def gen_stub(self, dest: Path, lines: Lines) -> None:
                 attr_type = "Any"
                 type_names.add(attr_type)
                 content.append(
-                    self.render(line, kind="attribute", name=name, attr_type=attr_type)
+                    self.render(line, kind=kind, name=name, attr_type=attr_type)
                 )
             else:
                 log.warning(

rst2pyi/tests/__init__.py

@@ -3,3 +3,5 @@
 
 from rst2pyi.core import Converter
 from rst2pyi.types import Line, Lines
+
+from .core import ConverterTest

rst2pyi/tests/core.py

@@ -0,0 +1,21 @@
+# Copyright 2019 John Reese
+# Licensed under the MIT license
+
+from unittest import TestCase
+
+from rst2pyi.core import Converter
+
+
+class ConverterTest(TestCase):
+    def test_split_types(self):
+        fn = Converter._split_types
+        self.assertEqual(fn("Optional[int]"), {"Optional", "int"})
+        self.assertEqual(fn("Union[int, str]"), {"Union", "int", "str"})
+        self.assertEqual(
+            fn("Optional[Union[int, str]]"), {"Optional", "Union", "int", "str"}
+        )
+
+    def test_convert_types(self):
+        fn = Converter._convert_type
+        self.assertEqual(fn("int or None"), "Optional[int]")
+        self.assertEqual(fn("int or str"), "Union[int, str]")

rst2pyi/types.py

@@ -20,6 +20,7 @@ class Config:
     )
     method_template: str = "    def {name}(self, {args}) -> {ret_type}: ..."
     attribute_template: str = "    {name}: {attr_type} = ..."
+    data_template: str = "\n# {source}:{lineno}\n{name}: {attr_type} = ..."
     function_template: str = (
         "\n# {source}:{lineno}\ndef {name}({args}) -> {ret_type}: ..."
     )