Refactor version detection logic

Author: fayazkhan

sqlalchemy_mptt/events.py

@@ -19,7 +19,7 @@
 from sqlalchemy.sql import func
 from sqlalchemy.orm.base import NO_VALUE
 
-from sqlalchemy_mptt.sqlalchemy_compat import case, select
+from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
 
 
 def _insert_subtree(
@@ -62,7 +62,7 @@ def _insert_subtree(
         .where(table.c.tree_id == parent_tree_id)
         .values(
             rgt=table.c.rgt + node_size,
-            lft=case(
+            lft=compat_layer.case(
                 (table.c.lft > left_sibling['lft'], table.c.lft + node_size),
                 else_=table.c.lft
             )
@@ -89,7 +89,7 @@ def mptt_before_insert(mapper, connection, instance):
         instance.right = 2
         instance.level = instance.get_default_level()
         tree_id = connection.scalar(
-            select(
+            compat_layer.select(
                 func.max(table.c.tree_id) + 1
             )
         ) or 1
@@ -99,7 +99,7 @@ def mptt_before_insert(mapper, connection, instance):
          parent_pos_right,
          parent_tree_id,
          parent_level) = connection.execute(
-            select(
+            compat_layer.select(
                 table.c.lft,
                 table.c.rgt,
                 table.c.tree_id,
@@ -115,11 +115,11 @@ def mptt_before_insert(mapper, connection, instance):
             .where(table.c.rgt >= parent_pos_right)
             .where(table.c.tree_id == parent_tree_id)
             .values(
-                lft=case(
+                lft=compat_layer.case(
                     (table.c.lft > parent_pos_right, table.c.lft + 2),
                     else_=table.c.lft
                 ),
-                rgt=case(
+                rgt=compat_layer.case(
                     (table.c.rgt >= parent_pos_right, table.c.rgt + 2),
                     else_=table.c.rgt
                 )
@@ -139,7 +139,7 @@ def mptt_before_delete(mapper, connection, instance, delete=True):
     db_pk = instance.get_pk_column()
     table_pk = getattr(table.c, db_pk.name)
     lft, rgt = connection.execute(
-        select(
+        compat_layer.select(
             table.c.lft,
             table.c.rgt
         ).where(
@@ -174,11 +174,11 @@ def mptt_before_delete(mapper, connection, instance, delete=True):
             .where(table.c.rgt > rgt)
             .where(table.c.tree_id == tree_id)
             .values(
-                lft=case(
+                lft=compat_layer.case(
                     (table.c.lft > lft, table.c.lft - delta),
                     else_=table.c.lft
                 ),
-                rgt=case(
+                rgt=compat_layer.case(
                     (table.c.rgt >= rgt, table.c.rgt - delta),
                     else_=table.c.rgt
                 )
@@ -210,7 +210,7 @@ def mptt_before_update(mapper, connection, instance):
             right_sibling_level,
             right_sibling_tree_id
         ) = connection.execute(
-            select(
+            compat_layer.select(
                 table.c.lft,
                 table.c.rgt,
                 table.c.parent_id,
@@ -221,7 +221,7 @@ def mptt_before_update(mapper, connection, instance):
             )
         ).fetchone()
         current_lvl_nodes = connection.execute(
-            select(
+            compat_layer.select(
                 table.c.lft,
                 table.c.rgt,
                 table.c.parent_id,
@@ -259,7 +259,7 @@ def mptt_before_update(mapper, connection, instance):
             left_sibling_parent,
             left_sibling_tree_id
         ) = connection.execute(
-            select(
+            compat_layer.select(
                 table.c.lft,
                 table.c.rgt,
                 table.c.parent_id,
@@ -282,7 +282,7 @@ def mptt_before_update(mapper, connection, instance):
         ORDER BY left_key
     """
     subtree = connection.execute(
-        select(table_pk)
+        compat_layer.select(table_pk)
         .where(
             and_(
                 table.c.lft >= instance.left,
@@ -306,7 +306,7 @@ def mptt_before_update(mapper, connection, instance):
         node_parent_id,
         node_level
     ) = connection.execute(
-        select(
+        compat_layer.select(
             table.c.lft,
             table.c.rgt,
             table.c.tree_id,
@@ -334,7 +334,7 @@ def mptt_before_update(mapper, connection, instance):
             parent_tree_id,
             parent_level
         ) = connection.execute(
-            select(
+            compat_layer.select(
                 table_pk,
                 table.c.rgt,
                 table.c.lft,
@@ -362,7 +362,7 @@ def mptt_before_update(mapper, connection, instance):
             parent_tree_id,
             parent_level
         ) = connection.execute(
-            select(
+            compat_layer.select(
                 table_pk,
                 table.c.rgt,
                 table.c.lft,
@@ -414,7 +414,7 @@ def mptt_before_update(mapper, connection, instance):
         # if just insert
         else:
             tree_id = connection.scalar(
-                select(
+                compat_layer.select(
                     func.max(table.c.tree_id) + 1
                 )
             )

sqlalchemy_mptt/sqlalchemy_compat.py

@@ -7,34 +7,49 @@
 import sqlalchemy as sa
 
 
-if sa.__version__ < '1.4':
-    from sqlalchemy.ext.declarative import declarative_base
-else:
-    from sqlalchemy.orm import declarative_base
+class LegacySQLAlchemyAPI:
+    """A class to provide compatibility for legacy SQLAlchemy versions (1.0 - 1.3)."""
 
+    @staticmethod
+    def declarative_base(*args, **kwargs):
+        from sqlalchemy.ext.declarative import declarative_base
+        return declarative_base(*args, **kwargs)
 
-def select(*args, **kwargs):
-    """Compatibility function for select."""
-    if sa.__version__ < '1.4':
+    @staticmethod
+    def select(*args, **kwargs):
         return sa.select(args, **kwargs)
-    else:
-        return sa.select(*args, **kwargs)
-
 
-def case(*args, **kwargs):
-    """Compatibility function for case."""
-    if sa.__version__ < '1.4':
+    @staticmethod
+    def case(*args, **kwargs):
         return sa.case(args, **kwargs)
-    else:
-        return sa.case(*args, **kwargs)
-
 
-def get(session, model, id):
-    """Compatibility function for getting an object by ID."""
-    if sa.__version__ < '1.4':
+    @staticmethod
+    def get(session, model, id):
         return session.query(model).get(id)
-    else:
+
+
+class ModernSQLAlchemyAPI:
+    """A class to provide compatibility for modern SQLAlchemy versions (1.4+)."""
+
+    @staticmethod
+    def declarative_base(*args, **kwargs):
+        from sqlalchemy.orm import declarative_base
+        return declarative_base(*args, **kwargs)
+
+    @staticmethod
+    def select(*args, **kwargs):
+        return sa.select(*args, **kwargs)
+
+    @staticmethod
+    def case(*args, **kwargs):
+        return sa.case(*args, **kwargs)
+
+    @staticmethod
+    def get(session, model, id):
         return session.get(model, id)
 
 
-__all__ = ["case", "declarative_base", "select"]
+if sa.__version__ < '1.4':
+    compat_layer = LegacySQLAlchemyAPI()
+else:
+    compat_layer = ModernSQLAlchemyAPI()

sqlalchemy_mptt/tests/test_events.py

@@ -19,11 +19,11 @@
 from sqlalchemy_mptt import mptt_sessionmaker
 
 from sqlalchemy_mptt.mixins import BaseNestedSets
-from sqlalchemy_mptt.sqlalchemy_compat import declarative_base
+from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
 from sqlalchemy_mptt.tests import TreeTestingMixin
 
 
-Base = declarative_base()
+Base = compat_layer.declarative_base()
 
 
 class Tree(Base, BaseNestedSets):

sqlalchemy_mptt/tests/test_inheritance.py

@@ -4,11 +4,11 @@
 from sqlalchemy.orm import sessionmaker
 
 from sqlalchemy_mptt.mixins import BaseNestedSets
-from sqlalchemy_mptt.sqlalchemy_compat import declarative_base, get
+from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
 from sqlalchemy_mptt.tests import TreeTestingMixin, failures_expected_on
 
 
-Base = declarative_base()
+Base = compat_layer.declarative_base()
 
 
 class GenericTree(Base, BaseNestedSets):
@@ -60,15 +60,15 @@ def test_create_generic(self):
         self.session.add(GenericTree(ppk=1))
         self.session.commit()
 
-        tree = get(self.session, GenericTree, 1)
+        tree = compat_layer.get(self.session, GenericTree, 1)
         self.assertEqual(tree.ppk, 1)
         self.assertEqual(tree.tree_id, 1)
 
     def test_create_spec(self):
         self.session.add(SpecializedTree(ppk=1))
         self.session.commit()
 
-        tree = get(self.session, SpecializedTree, 1)
+        tree = compat_layer.get(self.session, SpecializedTree, 1)
         self.assertEqual(tree.ppk, 1)
         self.assertEqual(tree.tree_id, 1)
 
@@ -84,21 +84,21 @@ def test_create_delete(self):
         self.session.add(parent)
         self.session.commit()
 
-        tree = get(self.session, SpecializedTree, 1)
+        tree = compat_layer.get(self.session, SpecializedTree, 1)
         self.assertEqual(tree.ppk, 1)
         self.assertEqual(tree.tree_id, 1)
 
         self.session.delete(child1)
         self.session.commit()
 
-        self.assertEqual(None, get(self.session, SpecializedTree, 2))
+        self.assertEqual(None, compat_layer.get(self.session, SpecializedTree, 2))
 
         self.session.delete(child2)
         self.session.commit()
 
-        self.assertEqual(None, get(self.session, SpecializedTree, 3))
-        self.assertEqual(None, get(self.session, SpecializedTree, 4))
-        self.assertEqual(None, get(self.session, SpecializedTree, 5))
+        self.assertEqual(None, compat_layer.get(self.session, SpecializedTree, 3))
+        self.assertEqual(None, compat_layer.get(self.session, SpecializedTree, 4))
+        self.assertEqual(None, compat_layer.get(self.session, SpecializedTree, 5))
 
 
 class TestGenericTree(TreeTestingMixin, unittest.TestCase):
@@ -116,7 +116,7 @@ def test_rebuild(self):
         super().test_rebuild()
 
 
-Base2 = declarative_base()
+Base2 = compat_layer.declarative_base()
 
 
 class BaseInheritance(Base2):

sqlalchemy_mptt/tests/test_mixins.py

@@ -14,10 +14,10 @@
 from sqlalchemy import Column, Integer
 
 from sqlalchemy_mptt.mixins import BaseNestedSets
-from sqlalchemy_mptt.sqlalchemy_compat import declarative_base
+from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
 
 
-Base = declarative_base()
+Base = compat_layer.declarative_base()
 
 
 class Tree2(Base, BaseNestedSets):

sqlalchemy_mptt/tests/test_stateful.py

@@ -11,10 +11,10 @@
 from sqlalchemy.orm import joinedload, sessionmaker
 
 from sqlalchemy_mptt import BaseNestedSets, mptt_sessionmaker
-from sqlalchemy_mptt.sqlalchemy_compat import declarative_base
+from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
 
 
-Base = declarative_base()
+Base = compat_layer.declarative_base()
 
 
 class Tree(Base, BaseNestedSets):