chore: refactor CRATE_META_OPTIONS from a dict to an Enum

Author: surister

cratedb_django/models/model.py

@@ -1,3 +1,5 @@
+from enum import Enum, auto
+
 from django.db import models, connection
 from django.db.models.base import ModelBase
 
@@ -6,28 +8,47 @@
 _OMITTED = type("OMITTED", (), {"__bool__": lambda _: False})
 OMITTED = _OMITTED()
 
-# dict of all the extra options a CrateDBModel Meta class has.
-# (name, default_value)
-CRATE_META_OPTIONS = {
-    "auto_refresh": False,  # Automatically refresh a table on inserts.
-    "partition_by": OMITTED,
-    "clustered_by": OMITTED,
-    "number_of_shards": OMITTED,
-}
+
+class CrateMetaOptions(Enum):
+    """
+    Represents the specific options a CrateDB table can have.
+    """
+
+    auto_refresh = auto(), False, False
+    partition_by = auto(), OMITTED, False
+    clustered_by = auto(), OMITTED, False
+    number_of_shards = auto(), OMITTED, False
+
+    def __init__(self, _, current_value, used_in_parameters_table):
+        self.current_value = current_value
+        self.used_in_parameters_table = used_in_parameters_table
+
+    @staticmethod
+    def options():
+        return list(CrateMetaOptions.__members__)
+
+    @classmethod
+    def by_name(cls, name):
+        for m in cls:
+            if m.value[1] == name:
+                return m
+        raise KeyError(name)
 
 
 class MetaCrate(ModelBase):
     def __new__(cls, name, bases, attrs, **kwargs):
-        crate_attrs = {}
+        _temp_crate_attrs = {}
 
         # todo document
 
         try:
             meta = attrs["Meta"]
-            for key, default_value in CRATE_META_OPTIONS.items():
-                crate_attrs[key] = getattr(meta, key, default_value)
-                if hasattr(meta, key):
-                    delattr(meta, key)
+            for option in CrateMetaOptions:
+                _temp_crate_attrs[option.name] = getattr(
+                    meta, option.name, option.current_value
+                )
+                if hasattr(meta, option.name):
+                    delattr(meta, option.name)
         except KeyError:
             # Has no meta class
             pass
@@ -36,7 +57,7 @@ def __new__(cls, name, bases, attrs, **kwargs):
 
         # Return back the crate_attrs we took from meta to the already
         # created object.
-        for k, v in crate_attrs.items():
+        for k, v in _temp_crate_attrs.items():
             setattr(o._meta, k, v)
         return o
 

tests/test_model.py

@@ -1,7 +1,7 @@
 import pytest
 
 from cratedb_django.models import CrateDBModel
-from cratedb_django.models.model import CRATE_META_OPTIONS, OMITTED
+from cratedb_django.models.model import CrateMetaOptions, OMITTED
 from cratedb_django import fields
 
 from django.forms.models import model_to_dict
@@ -57,6 +57,7 @@ def test_insert_model_field():
 
 
 def test_update_model():
+    assert SimpleModel.objects.count() == 0
     obj = SimpleModel.objects.create(field="text")
     SimpleModel.refresh()
     pk = obj.pk
@@ -131,15 +132,15 @@ class Meta:
             auto_refresh = True
 
     # Check all defaults are set.
-    for key, default_value in CRATE_META_OPTIONS.items():
-        assert key in NoMetaOptions._meta.__dict__
-        assert getattr(NoMetaOptions._meta, key) is default_value
+    for option in CrateMetaOptions:
+        assert option.name in NoMetaOptions._meta.__dict__
+        assert getattr(NoMetaOptions._meta, option.name) is option.current_value
 
     # Check the combination of user-defined + default.
     assert RefreshMetaOptions._meta.auto_refresh is True
     assert (
         RefreshMetaOptions._meta.partition_by
-        is CRATE_META_OPTIONS["partition_by"]
+        is CrateMetaOptions.partition_by.current_value
     )