1import datetime
  2
  3from plain.models import Model, models_registry
  4
  5
  6class JobParameter:
  7    """Base class for job parameter serialization/deserialization."""
  8
  9    STR_PREFIX = None  # Subclasses should define this
 10
 11    @classmethod
 12    def serialize(cls, value):
 13        """Return serialized string or None if can't handle this value."""
 14        return None
 15
 16    @classmethod
 17    def deserialize(cls, data):
 18        """Return deserialized value or None if can't handle this data."""
 19        return None
 20
 21    @classmethod
 22    def _extract_string_value(cls, data):
 23        """Extract value from string with prefix, return None if invalid format."""
 24        if not isinstance(data, str) or not cls.STR_PREFIX:
 25            return None
 26        if not data.startswith(cls.STR_PREFIX) or len(data) <= len(cls.STR_PREFIX):
 27            return None
 28        return data[len(cls.STR_PREFIX) :]
 29
 30
 31class ModelParameter(JobParameter):
 32    """Handle Plain model instances using a new string format."""
 33
 34    STR_PREFIX = "__plain://model/"
 35
 36    @classmethod
 37    def serialize(cls, value):
 38        if isinstance(value, Model):
 39            return f"{cls.STR_PREFIX}{value._meta.package_label}/{value._meta.model_name}/{value.id}"
 40        return None
 41
 42    @classmethod
 43    def deserialize(cls, data):
 44        if value_part := cls._extract_string_value(data):
 45            try:
 46                parts = value_part.split("/")
 47                if len(parts) == 3 and all(parts):
 48                    package, model_name, obj_id = parts
 49                    model = models_registry.get_model(package, model_name)
 50                    return model.objects.get(id=obj_id)
 51            except (ValueError, Exception):
 52                pass
 53        return None
 54
 55
 56class DateParameter(JobParameter):
 57    """Handle date objects."""
 58
 59    STR_PREFIX = "__plain://date/"
 60
 61    @classmethod
 62    def serialize(cls, value):
 63        if isinstance(value, datetime.date) and not isinstance(
 64            value, datetime.datetime
 65        ):
 66            return f"{cls.STR_PREFIX}{value.isoformat()}"
 67        return None
 68
 69    @classmethod
 70    def deserialize(cls, data):
 71        if value_part := cls._extract_string_value(data):
 72            try:
 73                return datetime.date.fromisoformat(value_part)
 74            except ValueError:
 75                pass
 76        return None
 77
 78
 79class DateTimeParameter(JobParameter):
 80    """Handle datetime objects."""
 81
 82    STR_PREFIX = "__plain://datetime/"
 83
 84    @classmethod
 85    def serialize(cls, value):
 86        if isinstance(value, datetime.datetime):
 87            return f"{cls.STR_PREFIX}{value.isoformat()}"
 88        return None
 89
 90    @classmethod
 91    def deserialize(cls, data):
 92        if value_part := cls._extract_string_value(data):
 93            try:
 94                return datetime.datetime.fromisoformat(value_part)
 95            except ValueError:
 96                pass
 97        return None
 98
 99
100class LegacyModelParameter(JobParameter):
101    """Legacy model parameter handling for backwards compatibility."""
102
103    STR_PREFIX = "gid://"
104
105    @classmethod
106    def serialize(cls, value):
107        # Don't serialize new instances with legacy format
108        return None
109
110    @classmethod
111    def deserialize(cls, data):
112        if value_part := cls._extract_string_value(data):
113            try:
114                package, model, obj_id = value_part.split("/")
115                model = models_registry.get_model(package, model)
116                return model.objects.get(id=obj_id)
117            except (ValueError, Exception):
118                pass
119        return None
120
121
122# Registry of parameter types to check in order
123# The order matters - more specific types should come first
124# DateTimeParameter must come before DateParameter since datetime is a subclass of date
125# LegacyModelParameter is last since it only handles deserialization
126PARAMETER_TYPES = [
127    ModelParameter,
128    DateTimeParameter,
129    DateParameter,
130    LegacyModelParameter,
131]
132
133
134class JobParameters:
135    """
136    Main interface for serializing and deserializing job parameters.
137    Uses the registered parameter types to handle different value types.
138    """
139
140    @staticmethod
141    def to_json(args, kwargs):
142        serialized_args = []
143        for arg in args:
144            serialized = JobParameters._serialize_value(arg)
145            serialized_args.append(serialized)
146
147        serialized_kwargs = {}
148        for key, value in kwargs.items():
149            serialized = JobParameters._serialize_value(value)
150            serialized_kwargs[key] = serialized
151
152        return {"args": serialized_args, "kwargs": serialized_kwargs}
153
154    @staticmethod
155    def _serialize_value(value):
156        """Serialize a single value using the registered parameter types."""
157        # Try each parameter type to see if it can serialize this value
158        for param_type in PARAMETER_TYPES:
159            result = param_type.serialize(value)
160            if result is not None:
161                return result
162
163        # If no parameter type can handle it, return as-is
164        return value
165
166    @staticmethod
167    def from_json(data):
168        args = []
169        for arg in data["args"]:
170            deserialized = JobParameters._deserialize_value(arg)
171            args.append(deserialized)
172
173        kwargs = {}
174        for key, value in data["kwargs"].items():
175            deserialized = JobParameters._deserialize_value(value)
176            kwargs[key] = deserialized
177
178        return args, kwargs
179
180    @staticmethod
181    def _deserialize_value(value):
182        """Deserialize a single value using the registered parameter types."""
183        # Try each parameter type to see if it can deserialize this value
184        for param_type in PARAMETER_TYPES:
185            result = param_type.deserialize(value)
186            if result is not None:
187                return result
188
189        # If no parameter type can handle it, return as-is
190        return value