Skip to content

Instantly share code, notes, and snippets.

@dimaqq

dimaqq/test_main_type_hint.py

Last active August 30, 2024 06:30
Show Gist options
  • Save dimaqq/f308257af8c3c4ecbe6f84953eeaac90 to your computer and use it in GitHub Desktop.
Save dimaqq/f308257af8c3c4ecbe6f84953eeaac90 to your computer and use it in GitHub Desktop.
Download ZIP
Type tests using pyright
Raw
test_main_type_hint.py
from dataclasses import dataclass
from typing import Optional, Protocol, Type
import ops
class CallableWithCharmClassOnly(Protocol):
"""Encapsulate main function type for simple charms.
Supports:
- ops.main(SomeCharm)
- ops.main(charm_class=SomeCharm)
"""
def __call__(self, charm_class: Type[ops.charm.CharmBase]): ...
class CallableWithCharmClassAndStorageFlag(Protocol):
"""Encapsulate main function type for advanced charms.
Supports permutations of:
- ops.main(SomeCharm, False)
- ops.main(charm_class=SomeCharm, use_juju_for_storage=False)
"""
def __call__(
self, charm_class: Type[ops.charm.CharmBase], use_juju_for_storage: Optional[bool] = None
): ...
class CallableWithoutArguments(Protocol):
"""Bad charm code should be caught by type checker.
For example:
- ops.main()
"""
def __call__(self): ...
@dataclass
class MainCalls:
simple: CallableWithCharmClassOnly
full: CallableWithCharmClassAndStorageFlag
bad: CallableWithoutArguments
sink = MainCalls(None, None, None) # type: ignore
def top_level_import() -> None:
import ops
sink.full = ops.main
sink.full = ops.main.main
sink.simple = ops.main
sink.simple = ops.main.main
sink.bad = ops.main # type: ignore[assignment]
sink.bad = ops.main.main # type: ignore[assignment]
def submodule_import() -> None:
import ops.main
sink.full = ops.main # type: ignore # https://github.com/microsoft/pyright/issues/8830
sink.full = ops.main.main
sink.simple = ops.main # type: ignore # https://github.com/microsoft/pyright/issues/8830
sink.simple = ops.main.main
sink.bad = ops.main # type: ignore[assignment]
sink.bad = ops.main.main # type: ignore[assignment]
def import_from_top_level_module() -> None:
from ops import main
sink.full = main
sink.full = main.main
sink.simple = main
sink.simple = main.main
sink.bad = main # type: ignore[assignment]
sink.bad = main.main # type: ignore[assignment]
def import_from_submodule() -> None:
from ops.main import main
sink.full = main
sink.simple = main
sink.bad = main # type: ignore[assignment]
@dimaqq
Copy link
Author

dimaqq commented Aug 30, 2024

Caveats:

Supports main decorated by stdlib decorators, behaves weird for proxies like def foo(*args, **kw).

Works when type is correct.

Shows reasonable, albeit verbose error when type is wrong:

test_main_type_hint.py:68:19 - error: Cannot assign to attribute "simple" for class "MainCalls"
    Type "(x: type[CharmBase]) -> Literal[42]" is not assignable to type "(charm_class: type[CharmBase]) -> None"
      Parameter name mismatch: "charm_class" versus "x"
      Function return type "Literal[42]" is incompatible with type "None"
        "Literal[42]" is not assignable to "None"

Call proxies:

Swallows fully untyped functions:

def proxy(*_a, **kw):  # type: ignore
    return ops.main(**kw)  # type: ignore

sink.simple = proxy  # pyright doesn't complain here

Complains correctly if some types are provided:

def proxy(*_a: List[Any], **kw: Dict[str, Union[bool, None]]):
    return ops.main(**kw)

sink.simple = proxy  # not assignable to type (pyright) / incompatible VarArg... (mypy)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment