Merge pull request #166 from arunsureshkumar/support-point-field-in-args

Support point field in args

Author: abawchen

graphene_mongo/advanced_types.py

@@ -3,7 +3,6 @@
 
 
 class FileFieldType(graphene.ObjectType):
-
     content_type = graphene.String()
     md5 = graphene.String()
     chunk_size = graphene.Int()
@@ -36,7 +35,6 @@ def resolve_data(self, info):
 
 
 class _CoordinatesTypeField(graphene.ObjectType):
-
     type = graphene.String()
 
     def resolve_type(self, info):
@@ -47,17 +45,19 @@ def resolve_coordinates(self, info):
 
 
 class PointFieldType(_CoordinatesTypeField):
-
     coordinates = graphene.List(graphene.Float)
 
 
-class PolygonFieldType(_CoordinatesTypeField):
+class PointFieldInputType(graphene.InputObjectType):
+    type = graphene.String(default_value="Point")
+    coordinates = graphene.List(graphene.Float, required=True)
 
+
+class PolygonFieldType(_CoordinatesTypeField):
     coordinates = graphene.List(graphene.List(graphene.List(graphene.Float)))
 
 
 class MultiPolygonFieldType(_CoordinatesTypeField):
-
     coordinates = graphene.List(
         graphene.List(graphene.List(graphene.List(graphene.Float)))
     )

graphene_mongo/converter.py

@@ -82,22 +82,22 @@ def convert_field_to_jsonstring(field, registry=None):
 
 @convert_mongoengine_field.register(mongoengine.PointField)
 def convert_point_to_field(field, registry=None):
-    return graphene.Field(advanced_types.PointFieldType)
+    return graphene.Field(advanced_types.PointFieldType, required=field.required)
 
 
 @convert_mongoengine_field.register(mongoengine.PolygonField)
 def convert_polygon_to_field(field, registry=None):
-    return graphene.Field(advanced_types.PolygonFieldType)
+    return graphene.Field(advanced_types.PolygonFieldType, required=field.required)
 
 
 @convert_mongoengine_field.register(mongoengine.MultiPolygonField)
 def convert_multipolygon_to_field(field, register=None):
-    return graphene.Field(advanced_types.MultiPolygonFieldType)
+    return graphene.Field(advanced_types.MultiPolygonFieldType, required=field.required)
 
 
 @convert_mongoengine_field.register(mongoengine.FileField)
 def convert_file_to_field(field, registry=None):
-    return graphene.Field(advanced_types.FileFieldType)
+    return graphene.Field(advanced_types.FileFieldType, required=field.required)
 
 
 @convert_mongoengine_field.register(mongoengine.ListField)
@@ -360,7 +360,7 @@ def dynamic_type():
             return None
         if isinstance(field, mongoengine.EmbeddedDocumentField):
             return graphene.Field(_type,
-                                  description=get_field_description(field, registry))
+                                  description=get_field_description(field, registry), required=field.required)
         field_resolver = None
         required = False
         if field.db_field is not None:

graphene_mongo/fields.py

@@ -24,7 +24,7 @@
     FileFieldType,
     PointFieldType,
     MultiPolygonFieldType,
-    PolygonFieldType,
+    PolygonFieldType, PointFieldInputType,
 )
 from .converter import convert_mongoengine_field, MongoEngineConversionError
 from .registry import get_global_registry
@@ -79,7 +79,7 @@ def registry(self):
     def args(self):
         return to_arguments(
             self._base_args or OrderedDict(),
-            dict(dict(self.field_args, **self.reference_args), **self.filter_args),
+            dict(dict(self.field_args, **self.advance_args), **self.filter_args),
         )
 
     @args.setter
@@ -149,35 +149,42 @@ def filter_args(self):
         if self._type._meta.filter_fields:
             for field, filter_collection in self._type._meta.filter_fields.items():
                 for each in filter_collection:
-                    filter_type = getattr(
-                        graphene,
-                        str(self._type._meta.fields[field].type).replace("!", ""),
-                    )
-
+                    if str(self._type._meta.fields[field].type) == 'PointFieldType':
+                        if each == 'max_distance':
+                            filter_type = graphene.Int
+                        else:
+                            filter_type = PointFieldInputType
+                    else:
+                        filter_type = getattr(
+                            graphene,
+                            str(self._type._meta.fields[field].type).replace("!", ""),
+                        )
                     # handle special cases
                     advanced_filter_types = {
                         "in": graphene.List(filter_type),
                         "nin": graphene.List(filter_type),
                         "all": graphene.List(filter_type),
                     }
-
                     filter_type = advanced_filter_types.get(each, filter_type)
                     filter_args[field + "__" + each] = graphene.Argument(
                         type=filter_type
                     )
-
         return filter_args
 
     @property
-    def reference_args(self):
-        def get_reference_field(r, kv):
+    def advance_args(self):
+        def get_advance_field(r, kv):
             field = kv[1]
             mongo_field = getattr(self.model, kv[0], None)
+            if isinstance(mongo_field, mongoengine.PointField):
+                r.update({kv[0]: graphene.Argument(PointFieldInputType)})
+                return r
             if isinstance(
                     mongo_field,
-                    (mongoengine.LazyReferenceField, mongoengine.ReferenceField),
+                    (mongoengine.LazyReferenceField, mongoengine.ReferenceField, mongoengine.GenericReferenceField),
             ):
-                field = convert_mongoengine_field(mongo_field, self.registry)
+                r.update({kv[0]: graphene.ID()})
+                return r
             if isinstance(mongo_field, mongoengine.GenericReferenceField):
                 r.update({kv[0]: graphene.ID()})
                 return r
@@ -192,7 +199,7 @@ def get_reference_field(r, kv):
 
             return r
 
-        return reduce(get_reference_field, self.fields.items(), {})
+        return reduce(get_advance_field, self.fields.items(), {})
 
     @property
     def fields(self):
@@ -220,6 +227,12 @@ def get_queryset(self, model, info, required_fields=list(), skip=None, limit=Non
                         reference_obj = get_document(arg["_cls"])(
                             pk=arg["_ref"].id)
                     hydrated_references[arg_name] = reference_obj
+                elif '__near' in arg_name and isinstance(getattr(self.model, arg_name.split('__')[0]),
+                                                         mongoengine.fields.PointField):
+                    location = args.pop(arg_name, None)
+                    hydrated_references[arg_name] = location["coordinates"]
+                    if (arg_name.split('__')[0] + "__max_distance") not in args:
+                        hydrated_references[arg_name.split('__')[0] + "__max_distance"] = 10000
                 elif arg_name == "id":
                     hydrated_references["id"] = from_global_id(args.pop("id", None))[1]
             args.update(hydrated_references)
@@ -381,10 +394,17 @@ def chained_resolver(self, resolver, is_partial, root, info, **args):
                                 self.filter_args.keys()):
                             args_copy.pop(arg_name)
                             if arg_name == '_id' and isinstance(arg, dict):
-                                args_copy['pk__in'] = arg['$in']
+                                operation = list(arg.keys())[0]
+                                args_copy['pk' + operation.replace('$', '__')] = arg[operation]
                             if '.' in arg_name:
                                 operation = list(arg.keys())[0]
                                 args_copy[arg_name.replace('.', '__') + operation.replace('$', '__')] = arg[operation]
+                        else:
+                            operations = ["$lte", "$gte", "$ne", "$in"]
+                            if isinstance(arg, dict) and any(op in arg for op in operations):
+                                operation = list(arg.keys())[0]
+                                args_copy[arg_name + operation.replace('$', '__')] = arg[operation]
+                                del args_copy[arg_name]
                     return self.default_resolver(root, info, required_fields, **args_copy)
                 else:
                     return resolved

graphene_mongo/tests/test_fields.py

@@ -9,7 +9,7 @@ def test_article_field_args():
     assert set(field.field_args.keys()) == field_args
 
     reference_args = {"editor", "reporter"}
-    assert set(field.reference_args.keys()) == reference_args
+    assert all(item in set(field.advance_args.keys()) for item in reference_args)
 
     default_args = {"after", "last", "first", "before"}
     args = field_args | reference_args | default_args