Add __all__ to unstable modules

This helps editors detect what can be imported from where, it at least improved my experience with VS Code
cognitedata · Nov 28, 2024 · b74bb4c · b74bb4c
1 parent d3b4a0a
commit b74bb4c
Show file tree

Hide file tree

Showing 7 changed files with 26 additions and 31 deletions.
diff --git a/cognite/extractorutils/unstable/configuration/loaders.py b/cognite/extractorutils/unstable/configuration/loaders.py
@@ -11,6 +11,9 @@
 from cognite.extractorutils.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
@@ -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
diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py
@@ -21,6 +21,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
@@ -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
@@ -16,6 +16,8 @@
 from ._messaging import RuntimeMessage
 from .base import ConfigRevision, ConfigType, Extractor
 
+__all__ = ["Runtime", "ExtractorType"]
+
 ExtractorType = TypeVar("ExtractorType", bound=Extractor)
 
 

diff --git 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):