From 62c5b5aa2322e5760715b46763c987faa116ff0e Mon Sep 17 00:00:00 2001
From: Mathias Lohne <mathias.lohne@cognite.com>
Date: Tue, 3 Dec 2024 11:37:22 +0100
Subject: [PATCH] Add `__all__` to unstable modules (#402)

This helps editors detect what can be imported from where, it at least
improved my experience with VS Code
---
 .../unstable/configuration/loaders.py         |  3 ++
 .../unstable/configuration/models.py          | 15 +++++++++
 .../extractorutils/unstable/core/__main__.py  | 31 -------------------
 cognite/extractorutils/unstable/core/base.py  |  2 ++
 .../extractorutils/unstable/core/errors.py    |  2 ++
 .../extractorutils/unstable/core/runtime.py   |  2 ++
 cognite/extractorutils/unstable/core/tasks.py |  2 ++
 7 files changed, 26 insertions(+), 31 deletions(-)
 delete mode 100644 cognite/extractorutils/unstable/core/__main__.py

diff --git a/cognite/extractorutils/unstable/configuration/loaders.py b/cognite/extractorutils/unstable/configuration/loaders.py
index e7c3f0a2..4a8313f6 100644
--- a/cognite/extractorutils/unstable/configuration/loaders.py
+++ b/cognite/extractorutils/unstable/configuration/loaders.py
@@ -12,6 +12,9 @@
 from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError
 from cognite.extractorutils.unstable.configuration.models import ConfigModel
 
+__all__ = ["ConfigFormat", "load_file", "load_from_cdf", "load_io", "load_dict"]
+
+
 _T = TypeVar("_T", bound=ConfigModel)
 
 
diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py
index 8791ebd6..d4373c84 100644
--- a/cognite/extractorutils/unstable/configuration/models.py
+++ b/cognite/extractorutils/unstable/configuration/models.py
@@ -19,6 +19,21 @@
 from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
 
+__all__ = [
+    "ConfigModel",
+    "AuthenticationConfig",
+    "TimeIntervalConfig",
+    "ConnectionConfig",
+    "CronConfig",
+    "IntervalConfig",
+    "ScheduleConfig",
+    "LogLevel",
+    "LogFileHandlerConfig",
+    "LogConsoleHandlerConfig",
+    "LogHandlerConfig",
+    "ExtractorConfig",
+]
+
 
 class ConfigModel(BaseModel):
     model_config = ConfigDict(
diff --git a/cognite/extractorutils/unstable/core/__main__.py b/cognite/extractorutils/unstable/core/__main__.py
deleted file mode 100644
index b8e9d06b..00000000
--- a/cognite/extractorutils/unstable/core/__main__.py
+++ /dev/null
@@ -1,31 +0,0 @@
-"""
-Example of how you would build an extractor with the new base class
-"""
-
-from cognite.extractorutils.unstable.configuration.models import ExtractorConfig
-
-from .base import Extractor
-from .runtime import Runtime
-
-
-class MyConfig(ExtractorConfig):
-    parameter_one: int
-    parameter_two: str
-
-
-class MyExtractor(Extractor[MyConfig]):
-    NAME = "Test extractor"
-    EXTERNAL_ID = "test-extractor"
-    DESCRIPTION = "Test of the new runtime"
-    VERSION = "1.0.0"
-    CONFIG_TYPE = MyConfig
-
-    def run(self) -> None:
-        self.logger.info("Started!")
-        if not self.cancellation_token.wait(10):
-            raise ValueError("Oops")
-
-
-if __name__ == "__main__":
-    runtime = Runtime(MyExtractor)
-    runtime.run()
diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py
index 1d71d296..7f3ba12c 100644
--- a/cognite/extractorutils/unstable/core/base.py
+++ b/cognite/extractorutils/unstable/core/base.py
@@ -30,6 +30,8 @@
 from cognite.extractorutils.unstable.scheduling import TaskScheduler
 from cognite.extractorutils.util import now
 
+__all__ = ["ConfigType", "ConfigRevision", "Extractor"]
+
 ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)
 ConfigRevision = Union[Literal["local"], int]
 
diff --git a/cognite/extractorutils/unstable/core/errors.py b/cognite/extractorutils/unstable/core/errors.py
index 2df0172a..9669cfe5 100644
--- a/cognite/extractorutils/unstable/core/errors.py
+++ b/cognite/extractorutils/unstable/core/errors.py
@@ -8,6 +8,8 @@
 if typing.TYPE_CHECKING:
     from .base import Extractor
 
+__all__ = ["Error", "ErrorLevel"]
+
 
 class ErrorLevel(Enum):
     warning = "warning"
diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py
index 732c8f66..5fedfbdd 100644
--- a/cognite/extractorutils/unstable/core/runtime.py
+++ b/cognite/extractorutils/unstable/core/runtime.py
@@ -22,6 +22,8 @@
 from ._messaging import RuntimeMessage
 from .base import ConfigRevision, ConfigType, Extractor, FullConfig
 
+__all__ = ["Runtime", "ExtractorType"]
+
 ExtractorType = TypeVar("ExtractorType", bound=Extractor)
 
 
diff --git a/cognite/extractorutils/unstable/core/tasks.py b/cognite/extractorutils/unstable/core/tasks.py
index a87dd4a9..0c5020ef 100644
--- a/cognite/extractorutils/unstable/core/tasks.py
+++ b/cognite/extractorutils/unstable/core/tasks.py
@@ -4,6 +4,8 @@
 
 from cognite.extractorutils.unstable.configuration.models import ScheduleConfig
 
+__all__ = ["ScheduledTask", "ContinuousTask", "StartupTask", "Task"]
+
 
 @dataclass
 class _Task(ABC):