From fb4ce6c5ed26016409aff6d445e0a9f698d8c957 Mon Sep 17 00:00:00 2001 From: Teon Banek Date: Wed, 11 Mar 2020 14:36:09 +0100 Subject: [PATCH] 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 | 56 ++++++++++++++++++++++++++-------------- 1 file changed, 37 insertions(+), 19 deletions(-) diff --git a/query_modules/example.py b/query_modules/example.py index 5040d3b49..4b9830de8 100644 --- a/query_modules/example.py +++ b/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) YIELD args, result; - Naturally, you may pass in different arguments or yield less fields. + CALL example.procedure(1, 2) YIELD args, vertex_count; + MATCH (n) CALL example.procedure(n, 1) YIELD * RETURN *; + + 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 - return mgp.Record(args=args_copy, result='Hello World!') + # Multiple rows can be produced by returning an iterable of mgp.Record. + return mgp.Record(args=args_copy, vertex_count=vertex_count, + avg_degree=avg_degree, props=props)