add: add support for overloaded functions (#196)

Signed-off-by: Stephen Whitlock <stephen@jshwisolutions.com>

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 [Unreleased](https://github.com/jshwi/docsig/compare/v0.35.0...HEAD)
 ------------------------------------------------------------------------
+### Added
+- add support for overloaded functions
 
 [0.35.0](https://github.com/jshwi/docsig/releases/tag/v0.35.0) - 2023-08-09
 ------------------------------------------------------------------------

docsig/_function.py

@@ -171,7 +171,7 @@ class _Signature(_DocSig):
     def __init__(  # pylint: disable=too-many-arguments
         self,
         arguments: _ast.Arguments,
-        returns: _ast.Module,
+        returns: _ast.Module | str,
         ismethod: bool = False,
         isstaticmethod: bool = False,
         ignore_args: bool = False,
@@ -197,7 +197,9 @@ def __init__(  # pylint: disable=too-many-arguments
             if a is not None and a.name
         )
 
-        self._rettype = self._get_rettype(returns)
+        self._rettype = (
+            returns if isinstance(returns, str) else self._get_rettype(returns)
+        )
         self._returns = str(self._rettype) != "None"
 
     def _get_rettype(self, returns: _ast.NodeNG | None) -> str | None:
@@ -315,6 +317,11 @@ def isproperty(self) -> bool:
         """Boolean value for whether function is a property."""
         return self.ismethod and self._decorated_with("property")
 
+    @property
+    def isoverloaded(self) -> bool:
+        """Boolean value for whether function is a property."""
+        return self._decorated_with("overload")
+
     @property
     def isinit(self) -> bool:
         """Boolean value for whether function is a class constructor."""

docsig/_module.py

@@ -32,9 +32,20 @@ def __init__(
         super().__init__()
         self._name = node.name
         self._path = f"{path}:" if path is not None else ""
+        overloads = []
+        returns = None
         for subnode in node.body:
             if isinstance(subnode, _ast.FunctionDef):
-                self.append(_Function(subnode, ignore_args, ignore_kwargs))
+                func = _Function(subnode, ignore_args, ignore_kwargs)
+                if func.isoverloaded:
+                    overloads.append(func.name)
+                    returns = func.signature.rettype
+                else:
+                    if func.name in overloads:
+                        subnode.returns = returns
+                        func = _Function(subnode, ignore_kwargs, ignore_kwargs)
+
+                    self.append(func)
 
     @property
     def path(self) -> str:

tests/__init__.py

@@ -6738,3 +6738,271 @@ def starmap(
     @property
     def expected(self) -> str:
         return ""
+
+
+@_templates.register
+class _MPassOverloadS(_BaseTemplate):
+    @property
+    def template(self) -> str:
+        return """
+from typing import overload
+
+
+@overload
+def process(response: None) -> None:
+    ...
+
+
+@overload
+def process(response: int) -> tuple[int, str]:
+    ...
+
+
+@overload
+def process(response: bytes) -> str:
+    ...
+
+
+def process(response):
+    \"\"\"process a response.
+
+    :param response: the response to process
+    :return: something depending on what the response is
+    \"\"\"
+"""
+
+    @property
+    def expected(self) -> str:
+        return ""
+
+
+@_templates.register
+class _MFailOverloadMissingReturnS(_BaseTemplate):
+    @property
+    def template(self) -> str:
+        return """
+from typing import overload
+
+
+@overload
+def process(response: None) -> None:
+    ...
+
+
+@overload
+def process(response: int) -> tuple[int, str]:
+    ...
+
+
+@overload
+def process(response: bytes) -> str:
+    ...
+
+
+def process(response):
+    \"\"\"process a response.
+
+    :param response: the response to process
+    \"\"\"
+"""
+
+    @property
+    def expected(self) -> str:
+        return f"""\
+module/file.py:20
+-----------------
+def process({CHECK}response) -> {CROSS}str:
+    \"\"\"
+    :param response: {CHECK}
+    :return: {CROSS}
+    \"\"\"
+
+{messages.E105}
+
+"""
+
+
+@_templates.register
+class _MFailOverloadMissingParamS(_BaseTemplate):
+    @property
+    def template(self) -> str:
+        return """
+from typing import overload
+
+
+@overload
+def process(response: None) -> None:
+    ...
+
+
+@overload
+def process(response: int) -> tuple[int, str]:
+    ...
+
+
+@overload
+def process(response: bytes) -> str:
+    ...
+
+
+def process(response):
+    \"\"\"process a response.
+
+    :return: something depending on what the response is
+    \"\"\"
+"""
+
+    @property
+    def expected(self) -> str:
+        return f"""\
+module/file.py:20
+-----------------
+def process({CROSS}response) -> {CHECK}str:
+    \"\"\"
+    :param None: {CROSS}
+    :return: {CHECK}
+    \"\"\"
+
+{messages.E103}
+
+"""
+
+
+@_templates.register
+class _MPassOverloadNoReturnS(_BaseTemplate):
+    @property
+    def template(self) -> str:
+        return """
+from typing import overload
+
+
+@overload
+def process(response: None) -> None:
+    ...
+
+
+@overload
+def process(response: int) -> None:
+    ...
+
+
+@overload
+def process(response: bytes) -> None:
+    ...
+
+
+def process(response):
+    \"\"\"process a response.
+
+    :param response: the response to process
+    \"\"\"
+"""
+
+    @property
+    def expected(self) -> str:
+        return ""
+
+
+@_templates.register
+class _MPassMultiOverloadsS(_BaseTemplate):
+    @property
+    def template(self) -> str:
+        return """
+from typing import overload
+
+
+@overload
+def process(response: None) -> None:
+    ...
+
+
+@overload
+def process(response: int) -> tuple[int, str]:
+    ...
+
+
+@overload
+def process(response: bytes) -> str:
+    ...
+
+
+def process(response):
+    \"\"\"process a response.
+
+    :param response: the response to process
+    :return: something depending on what the response is
+    \"\"\"
+
+@overload
+def another_process(response: int) -> tuple[int, str]:
+    ...
+
+
+@overload
+def another_process(response: bool) -> None:
+    ...
+
+
+@overload
+def another_process(response: str) -> int:
+    ...
+
+
+def another_process(response):
+    \"\"\"process another response.
+
+    :param response: the response to process
+    :return: something depending on what the response is
+    \"\"\"
+"""
+
+    @property
+    def expected(self) -> str:
+        return ""
+
+
+@_templates.register
+class _MFailOverloadNoReturnDocumentedS(_BaseTemplate):
+    @property
+    def template(self) -> str:
+        return """
+from typing import overload
+
+
+@overload
+def process(response: None) -> None:
+    ...
+
+
+@overload
+def process(response: int) -> None:
+    ...
+
+
+@overload
+def process(response: bytes) -> None:
+    ...
+
+
+def process(response):
+    \"\"\"process a response.
+
+    :param response: the response to process
+    :return: NoneType
+    \"\"\"
+"""
+
+    @property
+    def expected(self) -> str:
+        return f"""\
+module/file.py:20
+-----------------
+def process({CHECK}response) -> {CROSS}None:
+    \"\"\"
+    :param response: {CHECK}
+    :return: {CROSS}
+    \"\"\"
+
+{messages.E104}
+
+"""

whitelist.py

@@ -168,7 +168,13 @@
 _MFailG  # unused class (tests/__init__.py:5691)
 _MFailN  # unused class (tests/__init__.py:2211)
 _MFailNI  # unused class (tests/__init__.py:3982)
+_MFailOverloadMissingParamS  # unused class (tests/__init__.py:6825)
+_MFailOverloadMissingReturnS  # unused class (tests/__init__.py:6779)
+_MFailOverloadNoReturnDocumentedS  # unused class (tests/__init__.py:6964)
 _MFailS  # unused class (tests/__init__.py:611)
+_MPassMultiOverloadsS  # unused class (tests/__init__.py:6906)
+_MPassOverloadNoReturnS  # unused class (tests/__init__.py:6871)
+_MPassOverloadS  # unused class (tests/__init__.py:6743)
 _PBinOpG  # unused class (tests/__init__.py:5942)
 _PBinOpN  # unused class (tests/__init__.py:2516)
 _PBinOpNI  # unused class (tests/__init__.py:4270)