typing: Use @typing.overload instaed of Union type

Add method overload to explicitly declare str-to-str input/output, as christhompson commented in issue #32. Signed-off-by: ko-zu <[email protected]>
ko-zu · Jul 13, 2024 · 9aa4833 · 9aa4833
1 parent c35b563
commit 9aa4833
Showing 1 changed file with 61 additions and 1 deletion.
diff --git a/publicsuffixlist/__init__.py b/publicsuffixlist/__init__.py
@@ -9,7 +9,7 @@
 
 import os
 from collections.abc import Iterable as iterable, ByteString as bytestring
-from typing import Optional, Tuple, Union, Iterable, ByteString
+from typing import Optional, Tuple, Union, Iterable, ByteString, overload
 
 __all__ = ["PublicSuffixList"]
 
@@ -217,6 +217,18 @@ def _countpublic(self, labels, accept_unknown=None) -> int:
             return 1
         return 0
 
+    @overload
+    def suffix(self,
+               domain: str,
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[str]: ...
+    @overload
+    def suffix(self,
+               domain: Union[BytesTuple, Iterable[ByteString]],
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[BytesTuple]: ...
     def suffix(self,
                domain: RelaxDomain,
                accept_unknown: Optional[bool] = None,
@@ -225,6 +237,18 @@ def suffix(self,
         """ Alias for privatesuffix """
         return self.privatesuffix(domain, accept_unknown=accept_unknown, keep_case=keep_case)
 
+    @overload
+    def privatesuffix(self,
+               domain: str,
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[str]: ...
+    @overload
+    def privatesuffix(self,
+               domain: Union[BytesTuple, Iterable[ByteString]],
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[BytesTuple]: ...
     def privatesuffix(self,
                       domain: RelaxDomain,
                       accept_unknown: Optional[bool] = None,
@@ -249,6 +273,18 @@ def privatesuffix(self,
 
         return self._joinlabels(domain, labels, -(publen + 1), keep_case=keep_case)
 
+    @overload
+    def publicsuffix(self,
+               domain: str,
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[str]: ...
+    @overload
+    def publicsuffix(self,
+               domain: Union[BytesTuple, Iterable[ByteString]],
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[BytesTuple]: ...
     def publicsuffix(self,
                      domain: RelaxDomain,
                      accept_unknown: Optional[bool] = None,
@@ -285,6 +321,18 @@ def is_public(self, domain: RelaxDomain) -> bool:
         publen = self._countpublic(labels)
         return bool(publen and publen == len(labels))
 
+    @overload
+    def privateparts(self,
+               domain: str,
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[Tuple[str, ...]]: ...
+    @overload
+    def privateparts(self,
+               domain: Union[BytesTuple, Iterable[ByteString]],
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[Tuple[BytesTuple, ...]]: ...
     def privateparts(self,
                      domain: RelaxDomain,
                      *,
@@ -309,6 +357,18 @@ def privateparts(self,
             else:
                 return tuple(x.lower() for x in domain[:-(publen+1)]) + (priv,)
 
+    @overload
+    def subdomain(self,
+               domain: str,
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[str]: ...
+    @overload
+    def subdomain(self,
+               domain: Union[BytesTuple, Iterable[ByteString]],
+               accept_unknown: Optional[bool] = None,
+               *,
+               keep_case: bool = False) -> Optional[BytesTuple]: ...
     def subdomain(self,
                   domain: RelaxDomain,
                   depth: int,