fix: rework ops.main type hints to allow different import flavours

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)

test/test_main.py

@@ -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):
@@ -236,7 +236,7 @@ def __init__(self, framework: ops.Framework):
 
             with patch.dict(os.environ, fake_environ):
                 with patch('ops.main._emit_charm_event') as mock_charm_event:
-                    ops.main(MyCharm)  # type: ignore
+                    ops.main(MyCharm)
 
         assert mock_charm_event.call_count == 1
         return mock_charm_event.call_args[0][1]

test/test_main_invocations.py

@@ -0,0 +1,136 @@
+# 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 typing import Callable, Type
+from unittest.mock import Mock
+
+import pytest
+
+import ops
+
+Reset = Callable[[], None]
+
+
+def type_test_dummy(_arg: Callable[[Type[ops.CharmBase], bool], None]):
+    """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.
+    """
+
+
+@pytest.fixture
+def reset(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+    monkeypatch.setattr('sys.argv', ('hooks/install',))
+    monkeypatch.setattr('ops._main._emit_charm_event', Mock())
+    monkeypatch.setattr('ops._main._get_charm_dir', lambda: tmp_path)
+    monkeypatch.setattr('ops._main._Manager._setup_root_logging', Mock())
+    monkeypatch.setattr('ops.charm._evaluate_status', Mock())
+    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 (reset := lambda: os.environ.pop('OPERATOR_DISPATCH', None))
+
+    reset()
+
+
+class IdleCharm(ops.CharmBase):
+    def __init__(self, framework: ops.Framework):
+        super().__init__(framework)
+
+
+def test_top_level_import(reset: Reset):
+    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
+
+    ops.main(IdleCharm)
+
+    reset()
+    ops.main.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        ops.main()  # type: ignore
+
+    with pytest.raises(TypeError):
+        ops.main.main()  # type: ignore
+
+
+def test_submodule_import(reset: Reset):
+    import ops.main
+
+    type_test_dummy(ops.main.__call__)  # type: ignore FIXME
+    type_test_dummy(ops.main.main)
+    type_test_negative(ops.main.__call__)  # type: ignore
+    type_test_negative(ops.main.main)  # type: ignore
+
+    ops.main(IdleCharm)  # type: ignore FIXME
+
+    reset()
+    ops.main.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        ops.main()  # type: ignore
+
+    with pytest.raises(TypeError):
+        ops.main.main()  # type: ignore
+
+
+def test_import_from_top_level_module(reset: Reset):
+    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
+
+    main(IdleCharm)
+
+    reset()
+    main.main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        main()  # type: ignore
+
+    with pytest.raises(TypeError):
+        main.main()  # type: ignore
+
+
+def test_import_from_submodule(reset: Reset):
+    from ops.main import main
+
+    type_test_dummy(main)
+    type_test_negative(main)  # type: ignore
+
+    main(IdleCharm)
+
+    with pytest.raises(TypeError):
+        main()  # type: ignore