add overload stubs

Fixes: #1146
Signed-off-by: Vincent Fazio <vfazio@gmail.com>

diff --git a/alembic/context.pyi b/alembic/context.pyi
index 9871fadddb9a85b27adef371fe6c6c3c35360fdd..86345c4f6076239f58b1ba6e526917b2d159ab13 100644 (file)
--- a/alembic/context.pyi
+++ b/alembic/context.pyi
@@ -7,7 +7,9 @@ from typing import Callable
 from typing import ContextManager
 from typing import Dict
 from typing import List
+from typing import Literal
 from typing import Optional
+from typing import overload
 from typing import TextIO
 from typing import Tuple
 from typing import TYPE_CHECKING
@@ -644,8 +646,13 @@ def get_tag_argument() -> Optional[str]:
 
     """
 
+@overload
+def get_x_argument(as_dictionary: Literal[False]) -> List[str]: ...
+@overload
+def get_x_argument(as_dictionary: Literal[True]) -> Dict[str, str]: ...
+@overload
 def get_x_argument(
-    as_dictionary: bool = False,
+    as_dictionary: bool = ...,
 ) -> Union[List[str], Dict[str, str]]:
     """Return the value(s) passed for the ``-x`` argument, if any.