135 lines
4.1 KiB
Python
135 lines
4.1 KiB
Python
# Copyright 2023 Memgraph Ltd.
|
|
#
|
|
# Use of this software is governed by the Business Source License
|
|
# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
|
|
# License, and you may not use this file except in compliance with the Business Source License.
|
|
#
|
|
# As of the Change Date specified in that file, in accordance with
|
|
# the Business Source License, use of this software will be governed
|
|
# by the Apache License, Version 2.0, included in the file
|
|
# licenses/APL.txt.
|
|
|
|
import mgp
|
|
|
|
# isort: off
|
|
from common.shared import BaseClass, InitializationGraphMutable, InitializationUnderlyingGraphMutable
|
|
|
|
initialization_underlying_graph_mutable = InitializationUnderlyingGraphMutable()
|
|
|
|
|
|
def cleanup_underlying():
|
|
initialization_underlying_graph_mutable.reset()
|
|
|
|
|
|
def init_underlying_graph_is_mutable(ctx: mgp.ProcCtx, object: mgp.Any):
|
|
initialization_underlying_graph_mutable.set()
|
|
|
|
|
|
def underlying_graph_is_mutable(ctx: mgp.ProcCtx, object: mgp.Any) -> mgp.Record(mutable=bool, init_called=bool):
|
|
if initialization_underlying_graph_mutable.get_to_return() > 0:
|
|
initialization_underlying_graph_mutable.increment_returned(1)
|
|
return mgp.Record(
|
|
mutable=object.underlying_graph_is_mutable(), init_called=initialization_underlying_graph_mutable.get()
|
|
)
|
|
return []
|
|
|
|
|
|
# Register batched
|
|
mgp.add_batch_read_proc(underlying_graph_is_mutable, init_underlying_graph_is_mutable, cleanup_underlying)
|
|
|
|
|
|
initialization_graph_mutable = InitializationGraphMutable()
|
|
|
|
|
|
def init_graph_is_mutable(ctx: mgp.ProcCtx):
|
|
initialization_graph_mutable.set()
|
|
|
|
|
|
def graph_is_mutable(ctx: mgp.ProcCtx) -> mgp.Record(mutable=bool, init_called=bool):
|
|
if initialization_graph_mutable.get_to_return() > 0:
|
|
initialization_graph_mutable.increment_returned(1)
|
|
return mgp.Record(mutable=ctx.graph.is_mutable(), init_called=initialization_graph_mutable.get())
|
|
return []
|
|
|
|
|
|
def cleanup_graph():
|
|
initialization_graph_mutable.reset()
|
|
|
|
|
|
# Register batched
|
|
mgp.add_batch_read_proc(graph_is_mutable, init_graph_is_mutable, cleanup_graph)
|
|
|
|
|
|
class BatchingNums(BaseClass):
|
|
def __init__(self, nums_to_return):
|
|
super().__init__(nums_to_return)
|
|
self._nums = []
|
|
self._i = 0
|
|
|
|
|
|
batching_nums = BatchingNums(10)
|
|
|
|
|
|
def init_batching_nums(ctx: mgp.ProcCtx):
|
|
batching_nums.set()
|
|
batching_nums._nums = [i for i in range(1, 11)]
|
|
batching_nums._i = 0
|
|
|
|
|
|
def batch_nums(ctx: mgp.ProcCtx) -> mgp.Record(num=int, init_called=bool, is_valid=bool):
|
|
if batching_nums.get_to_return() > 0:
|
|
batching_nums.increment_returned(1)
|
|
batching_nums._i += 1
|
|
return mgp.Record(
|
|
num=batching_nums._nums[batching_nums._i - 1],
|
|
init_called=batching_nums.get(),
|
|
is_valid=ctx.graph.is_valid(),
|
|
)
|
|
return []
|
|
|
|
|
|
def cleanup_batching_nums():
|
|
batching_nums.reset()
|
|
batching_nums._i = 0
|
|
|
|
|
|
# Register batched
|
|
mgp.add_batch_read_proc(batch_nums, init_batching_nums, cleanup_batching_nums)
|
|
|
|
|
|
class BatchingVertices(BaseClass):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._vertices = []
|
|
self._i = 0
|
|
|
|
|
|
batching_vertices = BatchingVertices()
|
|
|
|
|
|
def init_batching_vertices(ctx: mgp.ProcCtx):
|
|
print("init called")
|
|
print("graph is mutable", ctx.graph.is_mutable())
|
|
batching_vertices.set()
|
|
batching_vertices._vertices = list(ctx.graph.vertices)
|
|
batching_vertices._i = 0
|
|
batching_vertices._num_to_return = len(batching_vertices._vertices)
|
|
|
|
|
|
def batch_vertices(ctx: mgp.ProcCtx) -> mgp.Record(vertex_id=int, init_called=bool):
|
|
if batching_vertices.get_to_return() == 0:
|
|
return []
|
|
batching_vertices.increment_returned(1)
|
|
return mgp.Record(vertex=batching_vertices._vertices[batching_vertices._i].id, init_called=batching_vertices.get())
|
|
|
|
|
|
def cleanup_batching_vertices():
|
|
batching_vertices.reset()
|
|
batching_vertices._vertices = []
|
|
batching_vertices._i = 0
|
|
batching_vertices._num_to_return = 0
|
|
|
|
|
|
# Register batched
|
|
mgp.add_batch_read_proc(batch_vertices, init_batching_vertices, cleanup_batching_vertices)
|