60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
from __future__ import absolute_import
|
|
|
|
from django.db.models.lookups import Exact
|
|
|
|
from bitfield.types import Bit, BitHandler
|
|
|
|
|
|
class BitQueryLookupWrapper(Exact): # NOQA
|
|
def process_lhs(self, compiler, connection, lhs=None):
|
|
lhs_sql, lhs_params = super(BitQueryLookupWrapper, self).process_lhs(
|
|
compiler, connection, lhs
|
|
)
|
|
|
|
if not isinstance(self.rhs, (BitHandler, Bit)):
|
|
return lhs_sql, lhs_params
|
|
|
|
op = " & " if self.rhs else " | "
|
|
rhs_sql, rhs_params = self.process_rhs(compiler, connection)
|
|
params = list(lhs_params)
|
|
params.extend(rhs_params)
|
|
|
|
return op.join((lhs_sql, rhs_sql)), params
|
|
|
|
def get_db_prep_lookup(self, value, connection):
|
|
v = value.mask if isinstance(value, (BitHandler, Bit)) else value
|
|
return super(BitQueryLookupWrapper, self).get_db_prep_lookup(v, connection)
|
|
|
|
def get_prep_lookup(self):
|
|
if isinstance(self.rhs, (BitHandler, Bit)):
|
|
return self.rhs # resolve at later stage, in get_db_prep_lookup
|
|
return super(BitQueryLookupWrapper, self).get_prep_lookup()
|
|
|
|
|
|
class BitQuerySaveWrapper(BitQueryLookupWrapper):
|
|
def as_sql(self, qn, connection):
|
|
"""
|
|
Create the proper SQL fragment. This inserts something like
|
|
"(T0.flags & value) != 0".
|
|
|
|
This will be called by Where.as_sql()
|
|
"""
|
|
engine = connection.settings_dict["ENGINE"].rsplit(".", -1)[-1]
|
|
if engine.startswith("postgres"):
|
|
XOR_OPERATOR = "#"
|
|
elif engine.startswith("sqlite"):
|
|
raise NotImplementedError
|
|
else:
|
|
XOR_OPERATOR = "^"
|
|
|
|
if self.bit:
|
|
return (
|
|
"%s.%s | %d" % (qn(self.table_alias), qn(self.column), self.bit.mask),
|
|
[],
|
|
)
|
|
return (
|
|
"%s.%s %s %d"
|
|
% (qn(self.table_alias), qn(self.column), XOR_OPERATOR, self.bit.mask),
|
|
[],
|
|
)
|