From 0c6bcb279e62f4c60a14f15990486092d877754b Mon Sep 17 00:00:00 2001
From: Sasha Rahlin <arahlin@uchicago.edu>
Date: Thu, 12 Sep 2024 17:36:09 -0500
Subject: [PATCH] Expose tuber.server.main() for user API

This PR allows using the main() function in the `__main__` block of user code.
For example:

```
registry = {"my_object": MyObject()}

if __name__ == "__main__":
    from tuber.server import main

    main(registry=registry)
```

The registry script then automatically provides the `tuberd` command-line
interface, with the registry pre-loaded by the user.

This PR also replaces the fragile `exec` call with an importlib interface for
loading a registry from a file path.
---
 tuber/server.py | 41 +++++++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 12 deletions(-)

diff --git a/tuber/server.py b/tuber/server.py
index 9c8a957..d0c0100 100644
--- a/tuber/server.py
+++ b/tuber/server.py
@@ -3,7 +3,7 @@
 import warnings
 from .codecs import Codecs
 
-__all__ = ["run"]
+__all__ = ["run", "main"]
 
 
 # request handling
@@ -370,20 +370,38 @@ def run(registry, json_module="json", port=80, webroot="/var/www/", max_age=3600
     )
 
 
-def main():
+def load_registry(filename):
     """
-    Server entry point
+    Load a user registry from a user provided python file
+    """
+
+    import importlib.util
+
+    spec = importlib.util.spec_from_file_location("tuber_registry", filename)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+
+    return mod.registry
+
+
+def main(registry=None):
+    """
+    Server entry point.
+
+    If supplied, run the server with the given registry.  Otherwise, use the
+    ``--registry`` command-line argument to provide a path to a registry file.
     """
 
     import argparse as ap
 
     P = ap.ArgumentParser(description="Tuber server")
-    P.add_argument(
-        "-r",
-        "--registry",
-        default="/usr/share/tuberd/registry.py",
-        help="Location of registry Python code",
-    )
+    if registry is None:
+        P.add_argument(
+            "-r",
+            "--registry",
+            default="/usr/share/tuberd/registry.py",
+            help="Location of registry Python code",
+        )
     P.add_argument(
         "-j",
         "--json",
@@ -409,9 +427,8 @@ def main():
     # setup environment
     os.environ["TUBER_SERVER"] = "1"
 
-    code = compile(open(args.registry, "r").read(), args.registry, "exec")
-    exec(code, globals())
-    args.registry = registry
+    # load registry
+    args.registry = registry if registry else load_registry(args.registry)
 
     run(**vars(args))