Update the example.py to use more API

Reviewers: ipaljak, dsantl, buda, tlastre

Reviewed By: ipaljak, buda

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D2718

query_modules/example.py

@@ -5,33 +5,51 @@ import mgp
 
 import copy
 
-# mgp.read_proc will use inspect.signature to get the arguments and types which
+
-# we map to our C API. Alternative registration examples are written below
-# this definition. Also, we can support other builtin types that make sense, as
-# well as the `typing` module, but `typing.Any` maps to `mgp.Nullable(mgp.Any)`
-# instead of `mgp.Any` which prevents `None`. Additionally, all `mgp.read_proc`
-# decorator registration relies on the assumption that they will have a
-# globally visible `mgp_module` on which procedures are registered. For this to
-# work, importing Python modules as Query Modules will need to go through a
-# single thread or should use a thread local variable. Currently, we load
-# everything on startup in a single thread, so there are no issues.
 @mgp.read_proc
 def procedure(context: mgp.ProcCtx,
               required_arg: mgp.Nullable[mgp.Any],
               optional_arg: mgp.Nullable[mgp.Any] = None
-              ) -> mgp.Record(result=str, args=list):
+              ) -> mgp.Record(args=list,
+                              vertex_count=int,
+                              avg_degree=mgp.Number,
+                              props=mgp.Nullable[mgp.Map]):
     '''
-    This example procedure returns 2 fields: `args` and `result`.
+    This example procedure returns 4 fields.
+
       * `args` is a copy of arguments passed to the procedure.
-      * `result` is the result of this procedure, a "Hello World!" string.
+      * `vertex_count` is the number of vertices in the database.
+      * `avg_degree` is the average degree of vertices.
+      * `props` is the properties map of the passed in `required_arg`, if it is
+        an Edge or a Vertex. In case of a Path instance, properties of the
+        starting vertex are returned.
+
     Any errors can be reported by raising an Exception.
 
     The procedure can be invoked in openCypher using the following calls:
-      CALL example.procedure(1, 2) YIELD args, result;
+      CALL example.procedure(1, 2) YIELD args, vertex_count;
-      CALL example.procedure(1) YIELD args, result;
+      MATCH (n) CALL example.procedure(n, 1) YIELD * RETURN *;
-    Naturally, you may pass in different arguments or yield less fields.
+
+    Naturally, you may pass in different arguments or yield different fields.
     '''
-    # Copying to make it equivalent to the C example.
+    # Create a properties map if we received an Edge, Vertex, or Path instance.
+    props = None
+    if isinstance(required_arg, (mgp.Edge, mgp.Vertex)):
+        props = dict(required_arg.properties.items())
+    elif isinstance(required_arg, mgp.Path):
+        start_vertex, = required_arg.vertices
+        props = dict(start_vertex.properties.items())
+    # Count the vertices and edges in the database; this may take a while.
+    vertex_count = 0
+    edge_count = 0
+    for v in context.graph.vertices:
+        vertex_count += 1
+        edge_count += sum(1 for e in v.in_edges)
+        edge_count += sum(1 for e in v.out_edges)
+    # Calculate the average degree, as if edges are not directed.
+    avg_degree = 0 if vertex_count == 0 else edge_count / vertex_count
+    # Copy the received arguments to make it equivalent to the C example.
     args_copy = [copy.deepcopy(required_arg), copy.deepcopy(optional_arg)]
-    # Multiple rows can be produced by returning an iterable of mgp.Record
+    # Multiple rows can be produced by returning an iterable of mgp.Record.
-    return mgp.Record(args=args_copy, result='Hello World!')
+    return mgp.Record(args=args_copy, vertex_count=vertex_count,
+                      avg_degree=avg_degree, props=props)