fix: rework ops.main type hints to allow different flavours (callable module)

ops/__init__.py

@@ -182,9 +182,7 @@
 # import was here previously
 from . import charm  # type: ignore # noqa: F401 `.charm` imported but unused
 
-# Import the main module, which we've overriden in main.py to be callable.
-# This allows "import ops; ops.main(Charm)" to work as expected.
-from . import main
+from . import main as _main
 
 # Explicitly import names from submodules so users can just "import ops" and
 # then use them as "ops.X".
@@ -321,3 +319,6 @@
 # rather than a runtime concern.
 
 from .version import version as __version__
+
+# This allows "import ops; ops.main(Charm)" to work as expected.
+main = _main._CallableMainModule('ops.main', _main.__doc__)

ops/main.py

@@ -26,6 +26,7 @@
 import sys
 import warnings
 from pathlib import Path
+from types import ModuleType
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 import ops.charm
@@ -555,14 +556,8 @@ def main(charm_class: Type[ops.charm.CharmBase], use_juju_for_storage: Optional[
         sys.exit(e.exit_code)
 
 
-# Make this module callable and call main(), so that "import ops" and then
-# "ops.main(Charm)" works as expected now that everything is imported in
-# ops/__init__.py. Idea from https://stackoverflow.com/a/48100440/68707
-class _CallableModule(sys.modules[__name__].__class__):
-    def __call__(
-        self, charm_class: Type[ops.charm.CharmBase], use_juju_for_storage: Optional[bool] = None
-    ):
-        return main(charm_class, use_juju_for_storage=use_juju_for_storage)
-
-
-sys.modules[__name__].__class__ = _CallableModule
+# Support old and new style main calls at run time and for type checking
+# - ops.main.main(SomeCharm)
+# - ops.main(SomeCharm)
+class _CallableMainModule(ModuleType):  # pyright: ignore[reportUnusedClass] as it's used in __init__.py
+    __call__ = main = staticmethod(main)

pyproject.toml

@@ -141,6 +141,12 @@ ignore = [
     # Manual dict comprehension.
     "PERF403",
 
+    # Convert Dict[X] to dict[x], which we cannot do while we support Python 3.8
+    "UP006",
+
+    # Convert Union[X, Y] to X | Y, which we cannot do while we support Python 3.8
+    "UP007",
+
     # Convert {} from `TypedDict` functional to class syntax
     # Note that since we have some `TypedDict`s that cannot use the class
     # syntax, we're currently choosing to be consistent in syntax even though

test/test_main.py

@@ -96,8 +96,8 @@ def __init__(
         self.check_name = check_name
 
 
-@patch('ops.main.setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
-@patch('ops.main._emit_charm_event', new=lambda *a, **kw: None)  # type: ignore
+@patch('ops._main.setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
+@patch('ops._main._emit_charm_event', new=lambda *a, **kw: None)  # type: ignore
 @patch('ops.charm._evaluate_status', new=lambda *a, **kw: None)  # type: ignore
 class TestCharmInit:
     @patch('sys.stderr', new_callable=io.StringIO)
@@ -148,7 +148,7 @@ def _check(
                 with fake_metadata.open('wb') as fh:
                     fh.write(b'name: test')
 
-                ops.main(charm_class, **kwargs)  # type: ignore
+                ops.main(charm_class, **kwargs)
 
     def test_init_signature_passthrough(self):
         class MyCharm(ops.CharmBase):
@@ -203,7 +203,7 @@ def test_controller_storage_deprecated(self):
 
 
 @patch('sys.argv', new=('hooks/config-changed',))
-@patch('ops.main._Manager._setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
+@patch('ops._main._Manager._setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
 @patch('ops.charm._evaluate_status', new=lambda *a, **kw: None)  # type: ignore
 class TestDispatch:
     def _check(self, *, with_dispatch: bool = False, dispatch_path: str = ''):
@@ -235,8 +235,8 @@ def __init__(self, framework: ops.Framework):
                 dispatch.chmod(0o755)
 
             with patch.dict(os.environ, fake_environ):
-                with patch('ops.main._emit_charm_event') as mock_charm_event:
-                    ops.main(MyCharm)  # type: ignore
+                with patch('ops._main._emit_charm_event') as mock_charm_event:
+                    ops.main(MyCharm)
 
         assert mock_charm_event.call_count == 1
         return mock_charm_event.call_args[0][1]

test/test_main_invocation.py

@@ -0,0 +1,107 @@
+# Copyright 2024 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from pathlib import Path
+from unittest.mock import Mock
+
+import pytest
+
+import ops
+
+
+@pytest.fixture
+def charm_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+    monkeypatch.setattr('sys.argv', ('hooks/install',))
+    monkeypatch.setattr('ops._main._emit_charm_event', Mock())
+    monkeypatch.setattr('ops._main._Manager._setup_root_logging', Mock())
+    monkeypatch.setattr('ops.charm._evaluate_status', Mock())
+    monkeypatch.setenv('JUJU_CHARM_DIR', str(tmp_path))
+    monkeypatch.setenv('JUJU_UNIT_NAME', 'test_main/0')
+    monkeypatch.setenv('JUJU_MODEL_NAME', 'mymodel')
+    monkeypatch.setenv('JUJU_DISPATCH_PATH', 'hooks/install')
+    monkeypatch.setenv('JUJU_VERSION', '3.5.0')
+    (tmp_path / 'metadata.yaml').write_text('name: test', encoding='utf-8')
+    (tmp_path / 'dispatch').absolute().touch(mode=0o755)
+
+    yield
+
+    os.environ.pop('OPERATOR_DISPATCH', None)
+
+
+class IdleCharm(ops.CharmBase):
+    def __init__(self, framework: ops.Framework):
+        super().__init__(framework)
+
+
+def test_top_level_import(charm_env: None):
+    import ops
+
+    ops.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        ops.main()  # type: ignore
+
+
+def test_top_level_import_legacy_call(charm_env: None):
+    import ops
+
+    ops.main.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        ops.main.main()  # type: ignore
+
+
+def test_submodule_import(charm_env: None):
+    import ops.main
+
+    ops.main(IdleCharm)  # type: ignore https://github.com/microsoft/pyright/issues/8830
+
+    with pytest.raises(TypeError):
+        ops.main()  # type: ignore
+
+
+def test_submodule_import_legacy_call(charm_env: None):
+    import ops.main
+
+    ops.main.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        ops.main.main()  # type: ignore
+
+
+def test_import_from_top_level_module(charm_env: None):
+    from ops import main
+
+    main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        main()  # type: ignore
+
+
+def test_import_from_top_level_module_legacy_call(charm_env: None):
+    from ops import main
+
+    main.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        main.main()  # type: ignore
+
+
+def test_legacy_import_from_submodule(charm_env: None):
+    from ops.main import main
+
+    main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        main()  # type: ignore

test/test_main_type_hint.py

@@ -0,0 +1,78 @@
+# Copyright 2024 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Callable, Type
+
+import ops
+
+
+def type_test_dummy(_arg: Callable[[Type[ops.CharmBase], bool], None]):
+    """
+    Helper to verify that
+    Usage:
+
+    from somewhere import main
+
+    type_test_dummy(main)
+    """
+
+
+def type_test_negative(_arg: Callable[[], None]):
+    """
+    Usage:
+
+    from somewhere import main
+
+    type_test_negative(main)  # type: ignore
+
+    The `reportUnnecessaryTypeIgnoreComment` setting is expected to kick up a fuss,
+    should the passed argument match the expected argument type.
+    """
+
+
+def top_level_import():
+    import ops
+
+    type_test_dummy(ops.main.__call__)  # pyright  is quirky
+    type_test_dummy(ops.main.main)
+
+    type_test_negative(ops.main.__call__)  # type: ignore
+    type_test_negative(ops.main.main)  # type: ignore
+
+
+def submodule_import():
+    import ops.main
+
+    type_test_dummy(ops.main.__call__)  # type: ignore https://github.com/microsoft/pyright/issues/8830
+    type_test_dummy(ops.main.main)
+
+    type_test_negative(ops.main.__call__)  # type: ignore
+    type_test_negative(ops.main.main)  # type: ignore
+
+
+def import_from_top_level_module():
+    from ops import main
+
+    type_test_dummy(main.__call__)
+    type_test_dummy(main.main)
+
+    type_test_negative(main.__call__)  # type: ignore
+    type_test_negative(main.main)  # type: ignore
+
+
+def import_from_submodule():
+    from ops.main import main
+
+    type_test_dummy(main)
+
+    type_test_negative(main)  # type: ignore