[unstable] Add File size Config

cognite/extractorutils/unstable/configuration/models.py

@@ -434,6 +434,148 @@ def _log_handler_default() -> list[LogHandlerConfig]:
     return [LogConsoleHandlerConfig(type="console", level=LogLevel.INFO)]
 
 
+class FileSizeConfig:
+    """
+    Configuration parameter for setting a file size.
+    """
+
+    def __init__(self, expression: str) -> None:
+        self._bytes, self._expression = FileSizeConfig._parse_expression(expression)
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:  # noqa: ANN401
+        """
+        Pydantic hook to define how this class should be serialized/deserialized.
+
+        This allows the class to be used as a field in Pydantic models.
+        """
+        return core_schema.no_info_after_validator_function(cls, handler(str | int))
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Two FileSizeConfig objects are equal if they have the same number of bytes.
+        """
+        if not isinstance(other, FileSizeConfig):
+            return NotImplemented
+        return self._bytes == other._bytes
+
+    def __hash__(self) -> int:
+        """
+        Hash function for FileSizeConfig based on the number of bytes.
+        """
+        return hash(self._bytes)
+
+    @classmethod
+    def _parse_expression(cls, expression: str) -> tuple[int, str]:
+        sizes = {
+            "kb": 1000,
+            "mb": 1_000_000,
+            "gb": 1_000_000_000,
+            "tb": 1_000_000_000_000,
+            "kib": 1024,
+            "mib": 1_048_576,
+            "gib": 1_073_741_824,
+            "tib": 1_099_511_627_776,
+        }
+        try:
+            return int(expression), expression
+        except ValueError:
+            pass
+        expression_lower = expression.lower().replace(" ", "")
+        for size in sizes:
+            if expression_lower.endswith(size):
+                num = float(expression_lower.replace(size, ""))
+                return int(num * sizes[size]), expression
+        raise InvalidConfigError(f"Invalid unit for file size: {expression}. Valid units: {sizes.keys()}")
+
+    @property
+    def bytes(self) -> int:
+        """
+        File size in bytes.
+        """
+        return self._bytes
+
+    @property
+    def kilobytes(self) -> float:
+        """
+        File size in kilobytes.
+        """
+        return self._bytes / 1000
+
+    @property
+    def megabytes(self) -> float:
+        """
+        File size in megabytes.
+        """
+        return self._bytes / 1_000_000
+
+    @property
+    def gigabytes(self) -> float:
+        """
+        File size in gigabytes.
+        """
+        return self._bytes / 1_000_000_000
+
+    @property
+    def terabytes(self) -> float:
+        """
+        File size in terabytes.
+        """
+        return self._bytes / 1_000_000_000_000
+
+    @property
+    def kibibytes(self) -> float:
+        """
+        File size in kibibytes (1024 bytes).
+        """
+        return self._bytes / 1024
+
+    @property
+    def mebibytes(self) -> float:
+        """
+        File size in mebibytes (1024 kibibytes).
+        """
+        return self._bytes / 1_048_576
+
+    @property
+    def gibibytes(self) -> float:
+        """
+        File size in gibibytes (1024 mebibytes).
+        """
+        return self._bytes / 1_073_741_824
+
+    @property
+    def tebibytes(self) -> float:
+        """
+        File size in tebibytes (1024 gibibytes).
+        """
+        return self._bytes / 1_099_511_627_776
+
+    def __int__(self) -> int:
+        """
+        Returns the file size as bytes.
+        """
+        return int(self._bytes)
+
+    def __float__(self) -> float:
+        """
+        Returns the file size as bytes.
+        """
+        return float(self._bytes)
+
+    def __str__(self) -> str:
+        """
+        Returns the file size as a human readable string.
+        """
+        return self._expression
+
+    def __repr__(self) -> str:
+        """
+        Returns the file size as a human readable string.
+        """
+        return self._expression
+
+
 class RawDestinationConfig(ConfigModel):
     """
     Configuration parameters for using Raw.

tests/test_unstable/test_configuration.py

@@ -1,12 +1,15 @@
 import os
+from dataclasses import field
 from io import StringIO
 
 import pytest
 
 from cognite.client.credentials import OAuthClientCredentials
 from cognite.extractorutils.unstable.configuration.loaders import ConfigFormat, load_io
 from cognite.extractorutils.unstable.configuration.models import (
+    ConfigModel,
     ConnectionConfig,
+    FileSizeConfig,
     TimeIntervalConfig,
     _ClientCredentialsConfig,
 )
@@ -212,3 +215,23 @@ def test_from_env() -> None:
 
     # Check that the produces cogniteclient object is valid
     assert len(client.assets.list(limit=1)) == 1
+
+
+class CustomFileConfig(ConfigModel):
+    file_size: FileSizeConfig = field(default_factory=lambda: FileSizeConfig("1MB"))
+    file_max_size: FileSizeConfig = field(default_factory=lambda: FileSizeConfig("10MiB"))
+
+
+def test_parse_file_size() -> None:
+    config_str = """
+file_size: 25MB
+file_max_size: 10MiB
+"""
+    stream = StringIO(config_str)
+    config = load_io(stream, ConfigFormat.YAML, CustomFileConfig)
+    assert config.file_size == FileSizeConfig("25MB")
+    assert config.file_size.bytes == 25_000_000
+    assert config.file_size._expression == "25MB"
+    assert config.file_max_size == FileSizeConfig("10MiB")
+    assert config.file_max_size.bytes == 10_485_760
+    assert config.file_max_size._expression == "10MiB"