feat(sdk): improved column definition class (#138)

sdk/radicalbit_platform_sdk/models/__init__.py

@@ -30,6 +30,7 @@
     RegressionModelQuality,
 )
 from .dataset_stats import DatasetStats
+from .field_type import FieldType
 from .file_upload_result import CurrentFileUpload, FileReference, ReferenceFileUpload
 from .job_status import JobStatus
 from .model_definition import (
@@ -39,6 +40,7 @@
     OutputType,
 )
 from .model_type import ModelType
+from .supported_types import SupportedTypes
 
 __all__ = [
     'OutputType',
@@ -76,4 +78,6 @@
     'CurrentFileUpload',
     'FileReference',
     'AwsCredentials',
+    'SupportedTypes',
+    'FieldType',
 ]

sdk/radicalbit_platform_sdk/models/column_definition.py

@@ -1,6 +1,13 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic.alias_generators import to_camel
 
+from radicalbit_platform_sdk.models.field_type import FieldType
+from radicalbit_platform_sdk.models.supported_types import SupportedTypes
 
-class ColumnDefinition(BaseModel):
+
+class ColumnDefinition(BaseModel, validate_assignment=True):
     name: str
-    type: str
+    type: SupportedTypes
+    field_type: FieldType = Field(alias='fieldType')
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)

sdk/radicalbit_platform_sdk/models/field_type.py

@@ -0,0 +1,23 @@
+from enum import Enum
+
+from radicalbit_platform_sdk.models.supported_types import SupportedTypes
+
+
+class FieldType(str, Enum):
+    categorical = 'categorical'
+    numerical = 'numerical'
+    datetime = 'datetime'
+
+    @staticmethod
+    def from_supported_type(value: SupportedTypes) -> 'FieldType':
+        match value:
+            case SupportedTypes.datetime:
+                return FieldType.datetime
+            case SupportedTypes.int:
+                return FieldType.numerical
+            case SupportedTypes.float:
+                return FieldType.numerical
+            case SupportedTypes.bool:
+                return FieldType.categorical
+            case SupportedTypes.string:
+                return FieldType.categorical

sdk/radicalbit_platform_sdk/models/supported_types.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class SupportedTypes(str, Enum):
+    string = 'string'
+    int = 'int'
+    float = 'float'
+    bool = 'bool'
+    datetime = 'datetime'