384 lines
13 KiB
Python
384 lines
13 KiB
Python
import json
|
|
|
|
from django.contrib.postgres import lookups
|
|
from django.contrib.postgres.forms import SimpleArrayField
|
|
from django.contrib.postgres.utils import (
|
|
CheckPostgresInstalledMixin,
|
|
prefix_validation_error,
|
|
)
|
|
from django.contrib.postgres.validators import ArrayMaxLengthValidator
|
|
from django.core import checks, exceptions
|
|
from django.db.models import Field, Func, IntegerField, Transform, Value
|
|
from django.db.models.fields.mixins import CheckFieldDefaultMixin
|
|
from django.db.models.lookups import Exact, In
|
|
from django.utils.translation import gettext_lazy as _
|
|
|
|
from .utils import AttributeSetter
|
|
|
|
__all__ = ["ArrayField"]
|
|
|
|
|
|
class ArrayField(CheckPostgresInstalledMixin, CheckFieldDefaultMixin, Field):
|
|
empty_strings_allowed = False
|
|
default_error_messages = {
|
|
"item_invalid": _("Item %(nth)s in the array did not validate:"),
|
|
"nested_array_mismatch": _("Nested arrays must have the same length."),
|
|
}
|
|
_default_hint = ("list", "[]")
|
|
|
|
def __init__(self, base_field, size=None, **kwargs):
|
|
self.base_field = base_field
|
|
self.db_collation = getattr(self.base_field, "db_collation", None)
|
|
self.size = size
|
|
if self.size:
|
|
self.default_validators = [
|
|
*self.default_validators,
|
|
ArrayMaxLengthValidator(self.size),
|
|
]
|
|
# For performance, only add a from_db_value() method if the base field
|
|
# implements it.
|
|
if hasattr(self.base_field, "from_db_value"):
|
|
self.from_db_value = self._from_db_value
|
|
super().__init__(**kwargs)
|
|
|
|
@property
|
|
def model(self):
|
|
try:
|
|
return self.__dict__["model"]
|
|
except KeyError:
|
|
raise AttributeError(
|
|
"'%s' object has no attribute 'model'" % self.__class__.__name__
|
|
)
|
|
|
|
@model.setter
|
|
def model(self, model):
|
|
self.__dict__["model"] = model
|
|
self.base_field.model = model
|
|
|
|
@classmethod
|
|
def _choices_is_value(cls, value):
|
|
return isinstance(value, (list, tuple)) or super()._choices_is_value(value)
|
|
|
|
def check(self, **kwargs):
|
|
errors = super().check(**kwargs)
|
|
if self.base_field.remote_field:
|
|
errors.append(
|
|
checks.Error(
|
|
"Base field for array cannot be a related field.",
|
|
obj=self,
|
|
id="postgres.E002",
|
|
)
|
|
)
|
|
else:
|
|
base_checks = self.base_field.check(**kwargs)
|
|
if base_checks:
|
|
error_messages = "\n ".join(
|
|
"%s (%s)" % (base_check.msg, base_check.id)
|
|
for base_check in base_checks
|
|
if isinstance(base_check, checks.Error)
|
|
# Prevent duplication of E005 in an E001 check.
|
|
and not base_check.id == "postgres.E005"
|
|
)
|
|
if error_messages:
|
|
errors.append(
|
|
checks.Error(
|
|
"Base field for array has errors:\n %s" % error_messages,
|
|
obj=self,
|
|
id="postgres.E001",
|
|
)
|
|
)
|
|
warning_messages = "\n ".join(
|
|
"%s (%s)" % (base_check.msg, base_check.id)
|
|
for base_check in base_checks
|
|
if isinstance(base_check, checks.Warning)
|
|
)
|
|
if warning_messages:
|
|
errors.append(
|
|
checks.Warning(
|
|
"Base field for array has warnings:\n %s"
|
|
% warning_messages,
|
|
obj=self,
|
|
id="postgres.W004",
|
|
)
|
|
)
|
|
return errors
|
|
|
|
def set_attributes_from_name(self, name):
|
|
super().set_attributes_from_name(name)
|
|
self.base_field.set_attributes_from_name(name)
|
|
|
|
@property
|
|
def description(self):
|
|
return "Array of %s" % self.base_field.description
|
|
|
|
def db_type(self, connection):
|
|
size = self.size or ""
|
|
return "%s[%s]" % (self.base_field.db_type(connection), size)
|
|
|
|
def cast_db_type(self, connection):
|
|
size = self.size or ""
|
|
return "%s[%s]" % (self.base_field.cast_db_type(connection), size)
|
|
|
|
def db_parameters(self, connection):
|
|
db_params = super().db_parameters(connection)
|
|
db_params["collation"] = self.db_collation
|
|
return db_params
|
|
|
|
def get_placeholder_sql(self, value, compiler, connection):
|
|
db_type = self.db_type(connection)
|
|
if hasattr(value, "as_sql"):
|
|
sql, params = compiler.compile(value)
|
|
return f"{sql}::{db_type}", params
|
|
return f"%s::{db_type}", (value,)
|
|
|
|
def get_db_prep_value(self, value, connection, prepared=False):
|
|
if isinstance(value, (list, tuple)):
|
|
return [
|
|
self.base_field.get_db_prep_value(i, connection, prepared=False)
|
|
for i in value
|
|
]
|
|
return value
|
|
|
|
def get_db_prep_save(self, value, connection):
|
|
if isinstance(value, (list, tuple)):
|
|
return [self.base_field.get_db_prep_save(i, connection) for i in value]
|
|
return value
|
|
|
|
def deconstruct(self):
|
|
name, path, args, kwargs = super().deconstruct()
|
|
if path == "django.contrib.postgres.fields.array.ArrayField":
|
|
path = "django.contrib.postgres.fields.ArrayField"
|
|
kwargs["base_field"] = self.base_field.clone()
|
|
if self.size is not None:
|
|
kwargs["size"] = self.size
|
|
return name, path, args, kwargs
|
|
|
|
def to_python(self, value):
|
|
if isinstance(value, str):
|
|
# Assume we're deserializing
|
|
vals = json.loads(value)
|
|
value = [self.base_field.to_python(val) for val in vals]
|
|
return value
|
|
|
|
def _from_db_value(self, value, expression, connection):
|
|
if value is None:
|
|
return value
|
|
return [
|
|
self.base_field.from_db_value(item, expression, connection)
|
|
for item in value
|
|
]
|
|
|
|
def value_to_string(self, obj):
|
|
values = []
|
|
vals = self.value_from_object(obj)
|
|
base_field = self.base_field
|
|
|
|
for val in vals:
|
|
if val is None:
|
|
values.append(None)
|
|
else:
|
|
obj = AttributeSetter(base_field.attname, val)
|
|
values.append(base_field.value_to_string(obj))
|
|
return json.dumps(values, ensure_ascii=False)
|
|
|
|
def get_transform(self, name):
|
|
transform = super().get_transform(name)
|
|
if transform:
|
|
return transform
|
|
if "_" not in name:
|
|
try:
|
|
index = int(name)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
index += 1 # postgres uses 1-indexing
|
|
return IndexTransformFactory(index, self.base_field)
|
|
try:
|
|
start, end = name.split("_")
|
|
start = int(start) + 1
|
|
end = int(end) # don't add one here because postgres slices are weird
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
return SliceTransformFactory(start, end)
|
|
|
|
def validate(self, value, model_instance):
|
|
super().validate(value, model_instance)
|
|
for index, part in enumerate(value):
|
|
try:
|
|
self.base_field.validate(part, model_instance)
|
|
except exceptions.ValidationError as error:
|
|
raise prefix_validation_error(
|
|
error,
|
|
prefix=self.error_messages["item_invalid"],
|
|
code="item_invalid",
|
|
params={"nth": index + 1},
|
|
)
|
|
if isinstance(self.base_field, ArrayField):
|
|
if len({len(i) for i in value}) > 1:
|
|
raise exceptions.ValidationError(
|
|
self.error_messages["nested_array_mismatch"],
|
|
code="nested_array_mismatch",
|
|
)
|
|
|
|
def run_validators(self, value):
|
|
super().run_validators(value)
|
|
for index, part in enumerate(value):
|
|
try:
|
|
self.base_field.run_validators(part)
|
|
except exceptions.ValidationError as error:
|
|
raise prefix_validation_error(
|
|
error,
|
|
prefix=self.error_messages["item_invalid"],
|
|
code="item_invalid",
|
|
params={"nth": index + 1},
|
|
)
|
|
|
|
def formfield(self, **kwargs):
|
|
return super().formfield(
|
|
**{
|
|
"form_class": SimpleArrayField,
|
|
"base_field": self.base_field.formfield(),
|
|
"max_length": self.size,
|
|
**kwargs,
|
|
}
|
|
)
|
|
|
|
def slice_expression(self, expression, start, length):
|
|
# If length is not provided, don't specify an end to slice to the end
|
|
# of the array.
|
|
end = None if length is None else start + length - 1
|
|
return SliceTransform(start, end, expression)
|
|
|
|
|
|
class ArrayRHSMixin:
|
|
def __init__(self, lhs, rhs):
|
|
# Don't wrap arrays that contains only None values, psycopg doesn't
|
|
# allow this.
|
|
if isinstance(rhs, (tuple, list)) and any(self._rhs_not_none_values(rhs)):
|
|
expressions = []
|
|
for value in rhs:
|
|
if not hasattr(value, "resolve_expression"):
|
|
field = lhs.output_field
|
|
value = Value(field.base_field.get_prep_value(value))
|
|
expressions.append(value)
|
|
rhs = Func(
|
|
*expressions,
|
|
function="ARRAY",
|
|
template="%(function)s[%(expressions)s]",
|
|
)
|
|
super().__init__(lhs, rhs)
|
|
|
|
def process_rhs(self, compiler, connection):
|
|
rhs, rhs_params = super().process_rhs(compiler, connection)
|
|
cast_type = self.lhs.output_field.cast_db_type(connection)
|
|
return "%s::%s" % (rhs, cast_type), rhs_params
|
|
|
|
def _rhs_not_none_values(self, rhs):
|
|
for x in rhs:
|
|
if isinstance(x, (list, tuple)):
|
|
yield from self._rhs_not_none_values(x)
|
|
elif x is not None:
|
|
yield True
|
|
|
|
|
|
@ArrayField.register_lookup
|
|
class ArrayContains(ArrayRHSMixin, lookups.DataContains):
|
|
pass
|
|
|
|
|
|
@ArrayField.register_lookup
|
|
class ArrayContainedBy(ArrayRHSMixin, lookups.ContainedBy):
|
|
pass
|
|
|
|
|
|
@ArrayField.register_lookup
|
|
class ArrayExact(ArrayRHSMixin, Exact):
|
|
pass
|
|
|
|
|
|
@ArrayField.register_lookup
|
|
class ArrayOverlap(ArrayRHSMixin, lookups.Overlap):
|
|
pass
|
|
|
|
|
|
@ArrayField.register_lookup
|
|
class ArrayLenTransform(Transform):
|
|
lookup_name = "len"
|
|
output_field = IntegerField()
|
|
|
|
def as_sql(self, compiler, connection):
|
|
lhs, params = compiler.compile(self.lhs)
|
|
# Distinguish NULL and empty arrays
|
|
return (
|
|
"CASE WHEN %(lhs)s IS NULL THEN NULL ELSE "
|
|
"coalesce(array_length(%(lhs)s, 1), 0) END"
|
|
) % {"lhs": lhs}, params * 2
|
|
|
|
|
|
@ArrayField.register_lookup
|
|
class ArrayInLookup(In):
|
|
def get_prep_lookup(self):
|
|
values = super().get_prep_lookup()
|
|
if hasattr(values, "resolve_expression"):
|
|
return values
|
|
# In.process_rhs() expects values to be hashable, so convert lists
|
|
# to tuples.
|
|
prepared_values = []
|
|
for value in values:
|
|
if hasattr(value, "resolve_expression"):
|
|
prepared_values.append(value)
|
|
else:
|
|
prepared_values.append(tuple(value))
|
|
return prepared_values
|
|
|
|
|
|
class IndexTransform(Transform):
|
|
def __init__(self, index, base_field, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.index = index
|
|
self.base_field = base_field
|
|
|
|
def as_sql(self, compiler, connection):
|
|
lhs, params = compiler.compile(self.lhs)
|
|
if not lhs.endswith("]"):
|
|
lhs = "(%s)" % lhs
|
|
return "%s[%%s]" % lhs, (*params, self.index)
|
|
|
|
@property
|
|
def output_field(self):
|
|
return self.base_field
|
|
|
|
|
|
class IndexTransformFactory:
|
|
def __init__(self, index, base_field):
|
|
self.index = index
|
|
self.base_field = base_field
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return IndexTransform(self.index, self.base_field, *args, **kwargs)
|
|
|
|
|
|
class SliceTransform(Transform):
|
|
def __init__(self, start, end, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.start = start
|
|
self.end = end
|
|
|
|
def as_sql(self, compiler, connection):
|
|
lhs, params = compiler.compile(self.lhs)
|
|
# self.start is set to 1 if slice start is not provided.
|
|
if self.end is None:
|
|
return f"({lhs})[%s:]", (*params, self.start)
|
|
else:
|
|
return f"({lhs})[%s:%s]", (*params, self.start, self.end)
|
|
|
|
|
|
class SliceTransformFactory:
|
|
def __init__(self, start, end):
|
|
self.start = start
|
|
self.end = end
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return SliceTransform(self.start, self.end, *args, **kwargs)
|