DEV-3250: Merge base classes

setup.cfg

@@ -45,8 +45,10 @@ where = src
 [options.extras_require]
 dev =
     coverage[toml]
+    more-itertools
     pytest
     pytest-cov>=2.7.1
+    typing-extensions
 lint =
     mypy
     pre-commit

src/psqlgraph/__init__.py

@@ -1,7 +1,44 @@
-from psqlgraph.base import create_all
+from psqlgraph import graph, voided
 from psqlgraph.edge import Edge, PolyEdge
+from psqlgraph.graph import AbstractEdge, AbstractNode
 from psqlgraph.node import Node, PolyNode
 from psqlgraph.psql import PsqlGraphDriver
-from psqlgraph.util import pg_property, sanitize, validate
-from psqlgraph.voided_edge import VoidedEdge
-from psqlgraph.voided_node import VoidedNode
+from psqlgraph.util import pg_property, sanitize
+from psqlgraph.voided import VoidedEdge, VoidedNode
+
+__all__ = (
+    "AbstractEdge",
+    "AbstractNode",
+    "Edge",
+    "Node",
+    "PolyEdge",
+    "PolyNode",
+    "PsqlGraphDriver",
+    "VoidedEdge",
+    "VoidedNode",
+    "create_all",
+    "drop_all",
+    "pg_property",
+    "sanitize",
+)
+
+
+def create_all(engine, base=graph.ORMBase):
+    """Create all tables associated with the provided declarative base, defaults to
+        ORMBase if not specified
+    Args:
+        engine (sqlalchemy.engine.Engine): active engine instance
+        base (sqlalchemy.ext.declarative.DeclarativeMeta): a declarative base class
+    """
+    base.metadata.create_all(engine)
+    voided.ORMBase.metadata.create_all(engine)
+
+
+def drop_all(engine, base=graph.ORMBase):
+    """Drops all tables associated with a given delcarative base
+    Args:
+        engine (sqlalchemy.engine.Engine): active engine instance
+        base (sqlalchemy.ext.declarative.DeclarativeMeta): a declarative base class
+    """
+    base.metadata.drop_all(engine)
+    voided.ORMBase.metadata.drop_all(engine)