From 37bb33cdbc1eae6a5c3cfc35eb472461b006287f Mon Sep 17 00:00:00 2001
From: karlicoss <karlicoss@gmail.com>
Date: Sat, 21 Oct 2023 22:25:16 +0100
Subject: [PATCH] experimental: add a hacky helper to import
 "original/shadowed" modules from within overlays

---
 my/core/experimental.py  | 64 ++++++++++++++++++++++++++++++++++++++++
 my/util/hpi_heartbeat.py | 54 +++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 my/core/experimental.py
 create mode 100644 my/util/hpi_heartbeat.py

diff --git a/my/core/experimental.py b/my/core/experimental.py
new file mode 100644
index 0000000..c10ba71
--- /dev/null
+++ b/my/core/experimental.py
@@ -0,0 +1,64 @@
+import sys
+from typing import Any, Dict, Optional
+import types
+
+
+# The idea behind this one is to support accessing "overlaid/shadowed" modules from namespace packages
+# See usage examples here:
+# - https://github.com/karlicoss/hpi-personal-overlay/blob/master/src/my/util/hpi_heartbeat.py
+# - https://github.com/karlicoss/hpi-personal-overlay/blob/master/src/my/twitter/all.py
+# Suppose you want to use my.twitter.talon, which isn't in the default all.py
+# You could just copy all.py to your personal overlay, but that would mean duplicating
+# all the code and possible upstream changes.
+# Alternatively, you could import the "original" my.twitter.all module from "overlay" my.twitter.all
+# _ORIG = import_original_module(__name__, __file__)
+# this would magically take care of package import path etc,
+# and should import the "original" my.twitter.all as _ORIG
+# After that you can call its methods, extend etc.
+def import_original_module(
+    module_name: str,
+    file: str,
+    *,
+    star: bool = False,
+    globals: Optional[Dict[str, Any]] = None,
+) -> types.ModuleType:
+    module_to_restore = sys.modules[module_name]
+
+    # NOTE: we really wanna to hack the actual package of the module
+    # rather than just top level my.
+    # since that would be a bit less disruptive
+    module_pkg = module_to_restore.__package__
+    assert module_pkg is not None
+    parent = sys.modules[module_pkg]
+
+    my_path = parent.__path__._path  # type: ignore[attr-defined]
+    my_path_orig = list(my_path)
+
+    def fixup_path() -> None:
+        for i, p in enumerate(my_path_orig):
+            starts = file.startswith(p)
+            if i == 0:
+                # not sure about this.. but I guess it'll always be 0th element?
+                assert starts, (my_path_orig, file)
+            if starts:
+                my_path.remove(p)
+        # should remove exactly one item
+        assert len(my_path) + 1 == len(my_path_orig), (my_path_orig, file)
+
+    try:
+        fixup_path()
+        try:
+            del sys.modules[module_name]
+            # NOTE: we're using __import__ instead of importlib.import_module
+            # since it's closer to the actual normal import (e.g. imports subpackages etc properly )
+            # fromlist=[None] forces it to return rightmost child
+            # (otherwise would just return 'my' package)
+            res = __import__(module_name, fromlist=[None])  # type: ignore[list-item]
+            if star:
+                assert globals is not None
+                globals.update({k: v for k, v in vars(res).items() if not k.startswith('_')})
+            return res
+        finally:
+            sys.modules[module_name] = module_to_restore
+    finally:
+        my_path[:] = my_path_orig
diff --git a/my/util/hpi_heartbeat.py b/my/util/hpi_heartbeat.py
new file mode 100644
index 0000000..84790a4
--- /dev/null
+++ b/my/util/hpi_heartbeat.py
@@ -0,0 +1,54 @@
+"""
+Just an helper module for testing HPI overlays
+In particular the behaviour of import_original_module function
+
+The idea of testing is that overlays extend this module, and add their own
+items to items(), and the checker asserts all overlays have contributed.
+"""
+from my.core import __NOT_HPI_MODULE__
+
+from dataclasses import dataclass
+from datetime import datetime
+import sys
+from typing import Iterator, List
+
+NOW = datetime.now()
+
+
+@dataclass
+class Item:
+    dt: datetime
+    message: str
+    path: List[str]
+
+
+def get_pkg_path() -> List[str]:
+    pkg = sys.modules[__package__]
+    return list(pkg.__path__)
+
+
+# NOTE: since we're hacking path for my.util
+# imports from my. should work as expected
+# (even though my.config is in the private config)
+from my.config import demo
+
+assert demo.username == 'todo'
+
+# however, this won't work while the module is imported
+# from my.util import extra
+# assert extra.message == 'EXTRA'
+# but it will work when we actually call the function (see below)
+
+
+def items() -> Iterator[Item]:
+    from my.config import demo
+
+    assert demo.username == 'todo'
+
+    # here the import works as expected, since by the time the function is called,
+    # all overlays were already processed and paths/sys.modules restored
+    from my.util import extra  # type: ignore[attr-defined]
+
+    assert extra.message == 'EXTRA'
+
+    yield Item(dt=NOW, message='hpi main', path=get_pkg_path())