Skip to content

Commit d88278a

Browse files
improved PipelineBuilder
1 parent dfd7e0a commit d88278a

6 files changed

Lines changed: 1084 additions & 469 deletions

File tree

mongoengine/base/document.py

Lines changed: 121 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1082,35 +1082,34 @@ def _lookup_field(cls, parts):
10821082
Returns:
10831083
A list of Field instances for fields that were found or
10841084
strings for sub-fields that weren't.
1085-
1086-
Example:
1087-
>>> user._lookup_field('name')
1088-
[<mongoengine.fields.StringField at 0x1119bff50>]
1089-
1090-
>>> user._lookup_field('roles')
1091-
[<mongoengine.fields.EmbeddedDocumentListField at 0x1119ec250>]
1092-
1093-
>>> user._lookup_field(['roles', 'role'])
1094-
[<mongoengine.fields.EmbeddedDocumentListField at 0x1119ec250>,
1095-
<mongoengine.fields.StringField at 0x1119ec050>]
1096-
1097-
>>> user._lookup_field('doesnt_exist')
1098-
raises LookUpError
1099-
1100-
>>> user._lookup_field(['roles', 'doesnt_exist'])
1101-
[<mongoengine.fields.EmbeddedDocumentListField at 0x1119ec250>,
1102-
'doesnt_exist']
1103-
11041085
"""
1105-
# TODO this method is WAY too complicated. Simplify it.
1106-
# TODO don't think returning a string for embedded non-existent fields is desired
1107-
11081086
ListField = _import_class("ListField")
11091087
DynamicField = _import_class("DynamicField")
1088+
DictField = _import_class("DictField")
1089+
MapField = _import_class("MapField")
1090+
ReferenceField = _import_class("ReferenceField")
1091+
GenericReferenceField = _import_class("GenericReferenceField")
11101092

11111093
if not isinstance(parts, (list, tuple)):
11121094
parts = [parts]
11131095

1096+
# Helper: resolve document classes for GenericReferenceField choices
1097+
def _resolve_generic_choices(generic_field):
1098+
from mongoengine.document import _DocumentRegistry
1099+
1100+
choices = getattr(generic_field, "choices", None) or ()
1101+
resolved = []
1102+
for ch in choices:
1103+
if isinstance(ch, str):
1104+
dc = _DocumentRegistry.get(ch)
1105+
elif isinstance(ch, type):
1106+
dc = _DocumentRegistry.get(ch.__name__)
1107+
else:
1108+
dc = None
1109+
if dc is not None:
1110+
resolved.append(dc)
1111+
return resolved
1112+
11141113
fields = []
11151114
field = None
11161115

@@ -1130,9 +1129,7 @@ def _lookup_field(cls, parts):
11301129
field = cls._fields[field_name]
11311130
elif cls._dynamic:
11321131
field = DynamicField(db_field=field_name)
1133-
elif cls._meta.get("allow_inheritance") or cls._meta.get(
1134-
"abstract", False
1135-
):
1132+
elif cls._meta.get("allow_inheritance") or cls._meta.get("abstract", False):
11361133
# 744: in case the field is defined in a subclass
11371134
for subcls in cls.__subclasses__():
11381135
try:
@@ -1146,55 +1143,117 @@ def _lookup_field(cls, parts):
11461143
raise LookUpError('Cannot resolve field "%s"' % field_name)
11471144
else:
11481145
raise LookUpError('Cannot resolve field "%s"' % field_name)
1149-
else:
1150-
ReferenceField = _import_class("ReferenceField")
1151-
GenericReferenceField = _import_class("GenericReferenceField")
11521146

1153-
# If previous field was a reference, throw an error (we
1154-
# cannot look up fields that are on references).
1155-
if isinstance(field, (ReferenceField, GenericReferenceField)):
1147+
fields.append(field)
1148+
continue
1149+
1150+
# ------------------------------------------------------------------
1151+
# JOINABLE PATH SUPPORT (ReferenceField / GenericReferenceField)
1152+
# plus ListField(ReferenceField/GenericReferenceField)
1153+
# ------------------------------------------------------------------
1154+
join_field = None
1155+
if isinstance(field, ReferenceField):
1156+
join_field = field
1157+
elif isinstance(field, GenericReferenceField):
1158+
join_field = field
1159+
elif isinstance(field, ListField) and isinstance(field.field, (ReferenceField, GenericReferenceField)):
1160+
join_field = field.field
1161+
1162+
if isinstance(join_field, ReferenceField):
1163+
target = getattr(join_field, "document_type", None) or getattr(join_field, "document_type_obj", None)
1164+
if target is None:
1165+
raise LookUpError('Cannot resolve reference target for "%s"' % join_field.name)
1166+
1167+
# Delegate resolution to referenced document. This does NOT perform a join;
1168+
# it only resolves the field definition so the aggregation/query layer can.
1169+
sub_field = target._lookup_field([field_name])[0]
1170+
field = sub_field
1171+
fields.append(field)
1172+
continue
1173+
1174+
if isinstance(join_field, GenericReferenceField):
1175+
# choices required in your design
1176+
choice_classes = _resolve_generic_choices(join_field)
1177+
if not choice_classes:
11561178
raise LookUpError(
1157-
"Cannot perform join in mongoDB: %s" % "__".join(parts)
1179+
'Cannot resolve GenericReferenceField choices for "%s"' % "__".join(parts)
11581180
)
11591181

1160-
# If the parent field has a "field" attribute which has a
1161-
# lookup_member method, call it to find the field
1162-
# corresponding to this iteration.
1163-
if hasattr(getattr(field, "field", None), "lookup_member"):
1164-
new_field = field.field.lookup_member(field_name)
1182+
resolved_fields = []
1183+
for dc in choice_classes:
1184+
resolved_fields.append(dc._lookup_field([field_name])[0])
11651185

1166-
# If the parent field is a DynamicField or if it's part of
1167-
# a DynamicDocument, mark current field as a DynamicField
1168-
# with db_name equal to the field name.
1169-
elif cls._dynamic and (
1170-
isinstance(field, DynamicField)
1171-
or getattr(getattr(field, "document_type", None), "_dynamic", None)
1172-
):
1173-
new_field = DynamicField(db_field=field_name)
1186+
# Must be consistent across choices (same Field class)
1187+
types = {type(f) for f in resolved_fields}
1188+
if len(types) != 1:
1189+
raise LookUpError(
1190+
'Ambiguous GenericReferenceField path "%s" (different field types across choices)'
1191+
% field_name
1192+
)
1193+
1194+
field = resolved_fields[0]
1195+
fields.append(field)
1196+
continue
11741197

1175-
# Else, try to use the parent field's lookup_member method
1176-
# to find the subfield.
1198+
# ------------------------------------------------------------------
1199+
# MapField/DictField key support:
1200+
# e.g. my_map__SOMEKEY__number
1201+
# SOMEKEY is a key, not a schema field.
1202+
# ------------------------------------------------------------------
1203+
if isinstance(field, (MapField, DictField)):
1204+
# Try normal resolution first (some containers expose lookup_member)
1205+
new_field = None
1206+
if hasattr(getattr(field, "field", None), "lookup_member"):
1207+
new_field = field.field.lookup_member(field_name)
11771208
elif hasattr(field, "lookup_member"):
11781209
new_field = field.lookup_member(field_name)
11791210

1180-
# Raise a LookUpError if all the other conditions failed.
1181-
else:
1182-
raise LookUpError(
1183-
"Cannot resolve subfield or operator {} "
1184-
"on the field {}".format(field_name, field.name)
1185-
)
1186-
1187-
# If current field still wasn't found and the parent field
1188-
# is a ComplexBaseField, add the name current field name and
1189-
# move on.
1190-
if not new_field and isinstance(field, ComplexBaseField):
1191-
fields.append(field_name)
1211+
if new_field:
1212+
field = new_field
1213+
fields.append(field)
11921214
continue
1193-
elif not new_field:
1194-
raise LookUpError('Cannot resolve field "%s"' % field_name)
11951215

1196-
field = new_field # update field to the new field type
1216+
# Treat as dictionary key token
1217+
fields.append(field_name)
1218+
# Descend into the container value field for the next segment
1219+
field = field.field
1220+
continue
1221+
1222+
# ------------------------------------------------------------------
1223+
# Original behavior for embedded/dynamic/complex fields
1224+
# ------------------------------------------------------------------
1225+
# If the parent field has a "field" attribute which has a
1226+
# lookup_member method, call it to find the field
1227+
if hasattr(getattr(field, "field", None), "lookup_member"):
1228+
new_field = field.field.lookup_member(field_name)
1229+
1230+
# If the parent field is a DynamicField or if it's part of
1231+
# a DynamicDocument, mark current field as a DynamicField
1232+
elif cls._dynamic and (
1233+
isinstance(field, DynamicField)
1234+
or getattr(getattr(field, "document_type", None), "_dynamic", None)
1235+
):
1236+
new_field = DynamicField(db_field=field_name)
1237+
1238+
# Else, try to use the parent field's lookup_member method
1239+
elif hasattr(field, "lookup_member"):
1240+
new_field = field.lookup_member(field_name)
1241+
1242+
else:
1243+
raise LookUpError(
1244+
"Cannot resolve subfield or operator {} "
1245+
"on the field {}".format(field_name, field.name)
1246+
)
1247+
1248+
# If current field still wasn't found and the parent field
1249+
# is a ComplexBaseField, add the name current field name and move on.
1250+
if not new_field and isinstance(field, ComplexBaseField):
1251+
fields.append(field_name)
1252+
continue
1253+
elif not new_field:
1254+
raise LookUpError('Cannot resolve field "%s"' % field_name)
11971255

1256+
field = new_field
11981257
fields.append(field)
11991258

12001259
return fields

0 commit comments

Comments
 (0)