From 626aceec35fa6e59afac7b47a5ede4ef4bb85fe9 Mon Sep 17 00:00:00 2001
From: Sveinung Gundersen <sveinugu@gmail.com>
Date: Thu, 9 Nov 2023 06:26:08 +0100
Subject: [PATCH] Start of code reorganisation. Unsure what the plan was, but
 it was at least related to organisation of the code to allow two types if
 reference documentation, one for primarily users of the library, and one for
 library developers.

---
 src/omnipy/api/__init__.py                  |  4 ++++
 src/omnipy/api/enums.py                     | 11 +++++++++++
 src/omnipy/api/exceptions.py                |  2 ++
 src/omnipy/api/protocols/__init__.py        |  1 +
 src/omnipy/api/protocols/public/__init__.py |  4 ++++
 src/omnipy/api/protocols/public/compute.py  |  4 ++++
 src/omnipy/api/typedefs.py                  |  2 ++
 src/omnipy/util/decorators.py               | 17 +++++++++++++++++
 8 files changed, 45 insertions(+)

diff --git a/src/omnipy/api/__init__.py b/src/omnipy/api/__init__.py
index e69de29b..b60acd97 100644
--- a/src/omnipy/api/__init__.py
+++ b/src/omnipy/api/__init__.py
@@ -0,0 +1,4 @@
+from .enums import *  # noqa: F401, F403
+from .exceptions import *  # noqa: F401, F403
+from .protocols import *  # noqa: F401, F403
+from .typedefs import *  # noqa: F401, F403
diff --git a/src/omnipy/api/enums.py b/src/omnipy/api/enums.py
index 3d0300aa..048ccdc5 100644
--- a/src/omnipy/api/enums.py
+++ b/src/omnipy/api/enums.py
@@ -1,5 +1,16 @@
 from enum import Enum, IntEnum
 
+__all__ = [
+    'PersistOutputsOptions',
+    'RestoreOutputsOptions',
+    'OutputStorageProtocolOptions',
+    'ConfigPersistOutputsOptions',
+    'ConfigRestoreOutputsOptions',
+    'ConfigOutputStorageProtocolOptions',
+    'EngineChoice',
+    'RunState',
+]
+
 
 class PersistOutputsOptions(str, Enum):
     """"""
diff --git a/src/omnipy/api/exceptions.py b/src/omnipy/api/exceptions.py
index f82768fe..fe9e58d7 100644
--- a/src/omnipy/api/exceptions.py
+++ b/src/omnipy/api/exceptions.py
@@ -1,5 +1,7 @@
 import omnipy.util.pydantic as pyd
 
+__all__ = ['JobStateException']
+
 
 class JobStateException(Exception):
     ...
diff --git a/src/omnipy/api/protocols/__init__.py b/src/omnipy/api/protocols/__init__.py
index e69de29b..15078f98 100644
--- a/src/omnipy/api/protocols/__init__.py
+++ b/src/omnipy/api/protocols/__init__.py
@@ -0,0 +1 @@
+from .public import *  # noqa: F401, F403
diff --git a/src/omnipy/api/protocols/public/__init__.py b/src/omnipy/api/protocols/public/__init__.py
index e69de29b..a55104b1 100644
--- a/src/omnipy/api/protocols/public/__init__.py
+++ b/src/omnipy/api/protocols/public/__init__.py
@@ -0,0 +1,4 @@
+from .compute import *  # noqa: F401, F403
+from .config import *  # noqa: F401, F403
+from .engine import *  # noqa: F401, F403
+from .hub import *  # noqa: F401, F403
diff --git a/src/omnipy/api/protocols/public/compute.py b/src/omnipy/api/protocols/public/compute.py
index 352d03dc..3d396a39 100644
--- a/src/omnipy/api/protocols/public/compute.py
+++ b/src/omnipy/api/protocols/public/compute.py
@@ -6,12 +6,16 @@
                                                       IsTaskTemplateArgsJob,
                                                       IsTaskTemplateArgsJobTemplate)
 from omnipy.api.protocols.private.compute.mixins import IsNestedContext
+from omnipy.util.decorators import export
+
+__all__ = []
 
 CallP = ParamSpec('CallP')
 RetT = TypeVar('RetT')
 RetCovT = TypeVar('RetCovT', covariant=True)
 
 
+@export
 class IsTaskTemplate(IsFuncArgJobTemplate['IsTaskTemplate[CallP, RetT]',
                                           'IsTask[CallP, RetT]',
                                           CallP,
diff --git a/src/omnipy/api/typedefs.py b/src/omnipy/api/typedefs.py
index 11313401..edcd4ffd 100644
--- a/src/omnipy/api/typedefs.py
+++ b/src/omnipy/api/typedefs.py
@@ -4,6 +4,8 @@
 from typing import _UnionGenericAlias  # type: ignore[attr-defined]
 from typing import _SpecialForm, Callable, ForwardRef, TypeAlias, TypeVar
 
+__all__: list[str] = []
+
 GeneralDecorator = Callable[[Callable], Callable]
 LocaleType: TypeAlias = str | tuple[str | None, str | None]
 
diff --git a/src/omnipy/util/decorators.py b/src/omnipy/util/decorators.py
index 849ed1c0..338f38e8 100644
--- a/src/omnipy/util/decorators.py
+++ b/src/omnipy/util/decorators.py
@@ -1,4 +1,5 @@
 from contextlib import AbstractContextManager
+import inspect
 from typing import Any, Callable, Concatenate, ContextManager, Generic, ParamSpec, TypeVar
 
 _DecoratedP = ParamSpec('_DecoratedP')
@@ -108,3 +109,19 @@ def __call__(self, arg: _ArgT) -> _ArgT:
                 return method(self._calling_obj_or_cls, arg)
 
     return SuperCaller
+
+
+"""
+Simple decorator for self-registering classes and functions as exported (in `__all__`).
+
+Adapted from https://stackoverflow.com/a/35710527
+"""
+
+
+def export(cls_or_fn):
+    module = inspect.getmodule(cls_or_fn)
+    if hasattr(module, '__all__'):
+        module.__all__.append(cls_or_fn.__name__)
+    else:
+        module.__all__ = [cls_or_fn.__name__]
+    return cls_or_fn