137 lines
5.7 KiB
Python
137 lines
5.7 KiB
Python
from django.db.models import QuerySet
|
|
from django.db.models.expressions import BaseExpression, Combinable
|
|
from django.db.models.query import ValuesIterable
|
|
from django.db.models.manager import Manager
|
|
from .adjuncts import BaseAdjunct
|
|
from .handlers import RecordDataclassHandler, RecordHandler
|
|
|
|
|
|
class RecordIterable(ValuesIterable):
|
|
"""
|
|
Iterable returned by records() and attached to it's queryset, that yields a record class for each row.
|
|
"""
|
|
|
|
def __iter__(self):
|
|
queryset = self.queryset
|
|
model = self.queryset.model
|
|
query = queryset.query
|
|
compiler = query.get_compiler(queryset.db)
|
|
record_data = getattr(queryset, '_record_kwargs', {})
|
|
try:
|
|
record_handler = queryset._record
|
|
except AttributeError:
|
|
raise AttributeError("The queryset lacks a _record entry, Is _clone copying that field on the QuerySet class?")
|
|
|
|
# extra(select=...) cols are always at the start of the row.
|
|
names = [
|
|
*query.extra_select,
|
|
*query.values_select,
|
|
*query.annotation_select,
|
|
]
|
|
indexes = range(len(names))
|
|
|
|
for row in compiler.results_iter(chunked_fetch=self.chunked_fetch, chunk_size=self.chunk_size):
|
|
dbdata = {names[i]: row[i] for i in indexes}
|
|
# post-processors will be able to rewrite the whole dictionary.
|
|
post_processors = []
|
|
# we overwrite db data bluntly for now. actually we would provide callbacks the current dict.
|
|
for k, v in record_data.items():
|
|
if v.resolves_field:
|
|
dbdata[k] = v.resolve(model, dbdata)
|
|
if v.post_processing:
|
|
post_processors.append(v)
|
|
if post_processors:
|
|
for processor in post_processors:
|
|
processed = processor.post_process(model, dbdata)
|
|
if processed is not None:
|
|
dbdata = processed
|
|
yield record_handler.create(**dbdata)
|
|
|
|
|
|
class RecordQuerySetMixin:
|
|
"""
|
|
Actual records() implementation.
|
|
As records() calls values(), the queryset is chain-cloned there.
|
|
|
|
You can mix this in into your QuerySet and Manager classes.
|
|
However make sure in QuerySet classes to implement _clone properly (or use RecordQuerySet instead)
|
|
"""
|
|
|
|
_record_handler = RecordDataclassHandler
|
|
|
|
def records(self, *args, **kwargs):
|
|
"""
|
|
generates record objects
|
|
|
|
Acts like values(), however:
|
|
- you can pass a record type or RecordHandler as first argument.
|
|
- if record type is not defined in records(), you have to define it on the queryset, or the model, with _default_record,
|
|
otherwise it will raise a RuntimeError.
|
|
- keyword arguments of type "Adjunct" are used as deferred values, and resolved independently.
|
|
- values() is called with every required_argument on the dataclass not handled by an Adjunct
|
|
"""
|
|
if len(args) and not isinstance(args[0], str):
|
|
# we assume this is our dataclass
|
|
# @TODO better checks.
|
|
handler = args[0]
|
|
args = args[1:]
|
|
else:
|
|
# determine dataclass.
|
|
handler = getattr(self, '_default_record', getattr(self.model, '_default_record', getattr(self, '_record', None)))
|
|
if not handler:
|
|
raise RuntimeError("Trying to records a class without destination class.")
|
|
|
|
if not isinstance(handler, RecordHandler):
|
|
handler = self._record_handler.wrap(handler)
|
|
|
|
all_keys = [*args, *kwargs.keys()]
|
|
unhandled_keys = list(set(handler.required_arguments) - set(all_keys))
|
|
args = [*args, *unhandled_keys]
|
|
|
|
# rebuild keyword arguments for values, by filtering out our adjuncts
|
|
new_kw = {}
|
|
adjuncts = {}
|
|
for k, v in kwargs.items():
|
|
if isinstance(v, BaseAdjunct):
|
|
# skip allows an adjunct to completely ignore a key.
|
|
if not v.skip:
|
|
adjuncts[k] = v
|
|
# check if we have to add to values. adjuncts can define a field to add here.
|
|
add_to_values = v.values_field()
|
|
if isinstance(add_to_values, str) and add_to_values not in args:
|
|
args.append(add_to_values)
|
|
elif isinstance(add_to_values, tuple):
|
|
new_kw[add_to_values[0]] = add_to_values[1]
|
|
elif isinstance(v, BaseExpression) or isinstance(v, Combinable) or hasattr(v, 'resolve_expression'):
|
|
new_kw[k] = v
|
|
else:
|
|
# this will fail in values() for now, but i do not want to hijack future django functionality here.
|
|
new_kw[k] = v
|
|
|
|
# copy ourself with values() and save the results on the cloned queryset values produces.
|
|
values = self.values(*args, **new_kw)
|
|
values._iterable_class = RecordIterable
|
|
values._record_kwargs = adjuncts
|
|
values._record = handler
|
|
return values
|
|
|
|
|
|
class RecordQuerySet(RecordQuerySetMixin, QuerySet):
|
|
# overwrite cloning. important.
|
|
def _clone(self):
|
|
c = super()._clone()
|
|
for key in ['_record', '_record_kwargs', '_record_handler', '_default_record']:
|
|
if hasattr(self, key):
|
|
setattr(c, key, getattr(self, key))
|
|
return c
|
|
|
|
# i use a mixin instead for better clarity with intellisense systems. records is completely safe, as it does not call _chain.
|
|
# however you can also simply do:
|
|
#class RecordManager(BaseManager.from_queryset(RecordQuerySet)):
|
|
# pass
|
|
|
|
|
|
class RecordManager(RecordQuerySetMixin, Manager):
|
|
def get_queryset(self):
|
|
return RecordQuerySet(self.model, using=self._db)
|