From b3aff0c62f98465363218c33075847749eea32fc Mon Sep 17 00:00:00 2001
From: Federico Baldassarre <baldassarre.fe@gmail.com>
Date: Tue, 29 Oct 2024 17:52:17 +0100
Subject: [PATCH] Enable automatically using import hook and env variable (#26)

---
 MANIFEST.in              |  2 ++
 _lovely_tensors_hook.pth |  1 +
 _lovely_tensors_hook.py  | 28 ++++++++++++++++++++++++++++
 setup.py                 | 26 ++++++++++++++++++++++++++
 4 files changed, 57 insertions(+)
 create mode 100644 _lovely_tensors_hook.pth
 create mode 100644 _lovely_tensors_hook.py

diff --git a/MANIFEST.in b/MANIFEST.in
index 5c0e7ce..0a43b48 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -2,4 +2,6 @@ include settings.ini
 include LICENSE
 include CONTRIBUTING.md
 include README.md
+include _lovely_tensors_hook.pth
+include _lovely_tensors_hook.py
 recursive-exclude * __pycache__
diff --git a/_lovely_tensors_hook.pth b/_lovely_tensors_hook.pth
new file mode 100644
index 0000000..84f51fd
--- /dev/null
+++ b/_lovely_tensors_hook.pth
@@ -0,0 +1 @@
+import _lovely_tensors_hook
diff --git a/_lovely_tensors_hook.py b/_lovely_tensors_hook.py
new file mode 100644
index 0000000..e77804d
--- /dev/null
+++ b/_lovely_tensors_hook.py
@@ -0,0 +1,28 @@
+"""
+This file should never be imported or executed by the user.
+
+It gets called by the site hook `_lovely_tensors_hook.pth` which in turn gets
+called automatically by python upon startup (https://docs.python.org/3/library/site.html)
+if the package lovely-tensors is installed.
+
+If the LOVEY_TENSORS environment variable is set to a "truthy" value, import and monkey patch.
+"""
+import os
+import sys
+
+
+if os.environ.get("LOVELY_TENSORS", "").lower() in {"1", "true", "yes"}:
+    try:
+        import lovely_tensors
+        lovely_tensors.monkey_patch()
+    except Exception as e:
+        ERROR_MESSAGE = """\
+Error: lovely_tensors.monkey_patch() failed with:
+
+{}
+
+If you uninstalled lovely_tensors, you should delete any '_lovely_tensors_hook.pth'
+file on your system and unset your 'LOVELY_TENSORS' environment variable.
+"""
+        print(ERROR_MESSAGE.format(e), file=sys.stderr)
+        raise
diff --git a/setup.py b/setup.py
index 6d52371..bd17475 100644
--- a/setup.py
+++ b/setup.py
@@ -1,8 +1,29 @@
 from pkg_resources import parse_version
 from configparser import ConfigParser
+import os.path
 import setuptools
+from setuptools.command.install_lib import install_lib
 assert parse_version(setuptools.__version__)>=parse_version('36.2')
 
+LOVELY_TENSORS_PTH = "_lovely_tensors_hook.pth"
+LOVELY_TENSORS_PY = "_lovely_tensors_hook.py"
+
+
+class InstallLibWithHook(install_lib):
+    def run(self):
+        install_lib.run(self)
+        outputs = []
+        for f in [LOVELY_TENSORS_PTH, LOVELY_TENSORS_PY]:
+            source = os.path.join(os.path.dirname(__file__), f)
+            dest = os.path.join(self.install_dir, f)
+            self.copy_file(source, dest)
+            outputs.append(dest)
+        self.outputs = outputs
+
+    def get_outputs(self):
+        return [*install_lib.get_outputs(self), *self.outputs]
+
+
 # note: all settings are in settings.ini; edit there, not here
 config = ConfigParser(delimiters=['='])
 config.read('settings.ini', encoding="utf-8")
@@ -52,6 +73,11 @@
         'console_scripts': cfg.get('console_scripts','').split(),
         'nbdev': [f'{cfg.get("lib_path")}={cfg.get("lib_path")}._modidx:d']
     },
+    # The import hooks mechanism is inspired by the one used in better_exceptions
+    # https://github.com/Qix-/better-exceptions/blob/f7f1476e57129dc74d241b4377b0df39c37bc8a7/setup.py
+    cmdclass = {
+        'install_lib': InstallLibWithHook,
+    },
     **setup_cfg)