def test_method_stub_should_include_params():
method = get_function("m",
params=["self", "i"],
ptypes=["int"],
rtype="None")
code = get_class("C", methods=[method])
assert get_stub(code) == "class C:\n def m(self, i: int) -> None: ...\n"
def test_if_method_decorated_callable_unknown_then_stub_should_ignore():
method = get_function("m",
params=["self"],
rtype="None",
decorators=["@foo()"])
code = get_class("C", methods=[method])
assert get_stub(code) == "class C:\n def m(self) -> None: ...\n"
def test_if_method_decorated_classmethod_then_stub_should_include_decorator():
method = get_function("m",
params=["cls"],
rtype="None",
decorators=["@classmethod"])
code = get_class("C", methods=[method])
assert get_stub(
code) == "class C:\n @classmethod\n def m(cls) -> None: ...\n"
def test_get_stub_typing_import_should_come_first():
code = "from x import A\n"
code += "\n\n" + get_function(
"f", params=["a", "l"], ptypes=["A", "List"], rtype="None")
assert (
get_stub(code) ==
"from typing import List\n\nfrom x import A\n\ndef f(a: A, l: List) -> None: ...\n"
)
def test_class_sig_should_not_overwrite_existing_init_sig():
method = get_function("__init__",
params=["self", "x"],
ptypes=["int"],
rtype="None")
code = get_class("C", sig="(str) -> int", methods=[method])
assert get_stub(
code) == "class C:\n def __init__(self, x: int) -> None: ...\n"
def test_if_method_decorated_property_then_stub_should_include_decorator():
method = get_function("m",
params=["self"],
rtype="None",
decorators=["@property"])
code = get_class("C", methods=[method])
assert get_stub(
code) == "class C:\n @property\n def m(self) -> None: ...\n"
def test_if_param_has_default_then_stub_should_include_ellipses():
code = get_function("f",
params=["i=0"],
ptypes=["Optional[int]"],
rtype="None")
assert (
get_stub(code) ==
"from typing import Optional\n\ndef f(i: Optional[int] = ...) -> None: ...\n"
)
def test_get_stub_comment_instance_variable():
method = get_function("m",
params=["self"],
rtype="None",
body="self.a = 42 # sig: int")
code = get_class("C", methods=[method])
assert (
get_stub(code) ==
"class C:\n a = ... # type: int\n def m(self) -> None: ...\n")
def test_stub_should_honor_kwonly_args_with_default():
code = get_function("f",
params=["i", "*", "j=0"],
ptypes=["int", "Optional[int]"],
rtype="None")
assert (
get_stub(code) ==
"from typing import Optional\n\ndef f(i: int, *, j: Optional[int] = ...) -> None: ...\n"
)
def test_if_param_type_multiple_typing_then_stub_should_include_multiple_import(
):
code = get_function("f",
params=["d"],
ptypes=["Dict[str, Any]"],
rtype="None")
assert (
get_stub(code) ==
"from typing import Any, Dict\n\ndef f(d: Dict[str, Any]) -> None: ...\n"
)
def test_if_returns_none_then_stub_should_return_none():
code = get_function("f", rtype="None")
assert get_stub(code) == "def f() -> None: ...\n"
def test_if_no_docstring_then_stub_should_be_empty():
code = get_function("f", desc="")
assert get_stub(code) == ""
def test_if_no_sig_then_stub_should_be_empty():
code = get_function("f")
assert get_stub(code) == ""
def test_stub_should_exclude_function_without_sig():
code = get_function("f", rtype="None")
code += "\n\n" + get_function("g", desc="")
assert get_stub(code) == "def f() -> None: ...\n"
def test_if_function_decorated_callable_unknown_then_stub_should_ignore():
code = get_function("f", rtype="None", decorators=["@foo()"])
assert get_stub(code) == "def f() -> None: ...\n"
def test_if_base_relative_imported_qualified_then_stub_should_include_import():
code = "from . import x\n"
code += "\n\n" + get_class("C", bases=["x.A"])
assert get_stub(code) == "from . import x\n\nclass C(x.A): ...\n"
def test_stub_should_use_alias_comment():
code = "# sigalias: B = int\n\nn = 42 # sig: B\n" ""
assert get_stub(code) == "B = int\n\nn = ... # type: B\n"
def test_module_variable_type_comment_relative_qualified_should_include_import(
):
code = "from . import x\n\nn = 42 # sig: x.A\n" ""
assert get_stub(code) == "from . import x\n\nn = ... # type: x.A\n"
def test_module_variable_type_comment_unimported_qualified_should_include_import(
):
code = "n = 42 # sig: x.y.A\n" ""
assert get_stub(code) == "import x.y\n\nn = ... # type: x.y.A\n"
def test_if_base_builtin_then_stub_should_include_base():
code = get_class("C", bases=["dict"])
assert get_stub(code) == "class C(dict): ...\n"
def test_module_variable_type_comment_from_imported_should_include_import():
code = "from x import A\n\nn = 42 # sig: A\n" ""
assert get_stub(code) == "from x import A\n\nn = ... # type: A\n"
def test_stub_should_exclude_unused_import():
code = "from x import A, B\n"
code += "\n\n" + get_function("f", rtype="A")
assert get_stub(code) == "from x import A\n\ndef f() -> A: ...\n"
def test_if_base_from_imported_qualified_then_stub_should_include_import():
code = "from x import y\n"
code += "\n\n" + get_class("C", bases=["y.A"])
assert get_stub(code) == "from x import y\n\nclass C(y.A): ...\n"
def test_if_base_unimported_qualified_then_stub_should_generate_import():
code = get_class("C", bases=["x.y.A"])
assert get_stub(code) == "import x.y\n\nclass C(x.y.A): ...\n"
def test_if_returns_builtin_then_stub_should_return_builtin():
code = get_function("f", rtype="int")
assert get_stub(code) == "def f() -> int: ...\n"
def test_stub_should_include_empty_class():
code = get_class("C")
assert get_stub(code) == "class C: ...\n"
def test_if_returns_from_imported_then_stub_should_include_import():
code = "from x import A\n"
code += "\n\n" + get_function("f", rtype="A")
assert get_stub(code) == "from x import A\n\ndef f() -> A: ...\n"
def test_class_sig_should_be_moved_to_init_method():
method = get_function("__init__", params=["self", "x"], rtype=None)
code = get_class("C", sig="(str) -> int", methods=[method])
assert get_stub(
code) == "class C:\n def __init__(self, x: str) -> int: ...\n"
def test_if_returns_imported_qualified_then_stub_should_include_import():
code = "import x\n"
code += "\n\n" + get_function("f", rtype="x.A")
assert get_stub(code) == "import x\n\ndef f() -> x.A: ...\n"
def test_module_variable_type_comment_builtin_should_be_ellipsized():
code = "n = 42 # sig: int\n"
assert get_stub(code) == "n = ... # type: int\n"