From bab9450441f7cc03e20b094b40721b951ff03be2 Mon Sep 17 00:00:00 2001
From: Leifur Halldor Asgeirsson <lasgeirsson@zerofail.com>
Date: Tue, 4 Oct 2016 12:12:12 -0400
Subject: [PATCH] implement app.register_from_module
---
goblin/app.py | 13 ++++++++++---
tests/register_models.py | 21 +++++++++++++++++++++
tests/test_app.py | 22 ++++++++++++++++++++++
3 files changed, 53 insertions(+), 3 deletions(-)
create mode 100644 tests/register_models.py
diff --git a/goblin/app.py b/goblin/app.py
index 89d05ca..c46d691 100644
--- a/goblin/app.py
+++ b/goblin/app.py
@@ -18,9 +18,9 @@
"""Goblin application class and class constructor"""
import collections
+import importlib
import logging
-from gremlin_python import process
from goblin import driver, element, session
@@ -120,8 +120,15 @@ class Goblin:
"""
self._cluster.config_from_json(filename)
- def register_from_module(self, modulename):
- raise NotImplementedError
+ def register_from_module(self, module, *, package=None):
+ if isinstance(module, str):
+ module = importlib.import_module(module, package)
+ elements = list()
+ for item_name in dir(module):
+ item = getattr(module, item_name)
+ if isinstance(item, element.ElementMeta):
+ elements.append(item)
+ self.register(*elements)
async def session(self, *, use_session=False, processor='', op='eval',
aliases=None):
diff --git a/tests/register_models.py b/tests/register_models.py
new file mode 100644
index 0000000..6a6eb78
--- /dev/null
+++ b/tests/register_models.py
@@ -0,0 +1,21 @@
+from goblin import element
+
+
+class TestRegisterVertex1(element.Vertex):
+ pass
+
+
+class TestRegisterVertex2(element.Vertex):
+ pass
+
+
+class TestRegisterEdge1(element.Edge):
+ pass
+
+
+class TestRegisterEdge2(element.Edge):
+ pass
+
+
+class NotAModelShouldNotBeRegistered:
+ pass
diff --git a/tests/test_app.py b/tests/test_app.py
index 89067f1..73f8237 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -17,11 +17,33 @@
import pytest
+import goblin
from goblin import element
from goblin.driver import serializer
from gremlin_python import process
+def test_register_from_module(app):
+ import register_models
+ app.register_from_module(register_models)
+ vertices, edges = app._vertices.values(), app._edges.values()
+ assert register_models.TestRegisterVertex1 in vertices
+ assert register_models.TestRegisterVertex2 in vertices
+ assert register_models.TestRegisterEdge1 in edges
+ assert register_models.TestRegisterEdge2 in edges
+
+
+def test_register_from_module_string(app):
+ app.register_from_module('register_models', package=__package__)
+ vertices, edges = app._vertices.values(), app._edges.values()
+
+ import register_models
+ assert register_models.TestRegisterVertex1 in vertices
+ assert register_models.TestRegisterVertex2 in vertices
+ assert register_models.TestRegisterEdge1 in edges
+ assert register_models.TestRegisterEdge2 in edges
+
+
@pytest.mark.asyncio
async def test_registry(app, person, place, knows, lives_in):
assert len(app.vertices) == 2
--
GitLab