Clean elements.py and fix tests in PyPy

geoalchemy · Mar 27, 2020 · 8605373 · 8605373
1 parent 585962c
commit 8605373
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/geoalchemy2/elements.py b/geoalchemy2/elements.py
@@ -45,10 +45,13 @@ def __init__(self, data, srid=-1, extended=False):
         self.data = data
         self.extended = extended
         if self.extended:
-            args = [self.geom_from_extended_version, self.data]
+            func_name = self.geom_from_extended_version
+            args = [self.data]
         else:
-            args = [self.geom_from, self.data, self.srid]
-        self.function_expr = getattr(functions.func, args[0])(*args[1:])
+            func_name = self.geom_from
+            args = [self.data, self.srid]
+
+        self.function_expr = getattr(functions.func, func_name)(*args)
 
     def __str__(self):
         return self.desc
@@ -96,10 +99,11 @@ def __setstate__(self, state):
         self.srid = state['srid']
         self.extended = state['extended']
         self.data = self._data_from_desc(state['data'])
-        args = [state['name'], self.data]
+        func_name = state['name']
+        args = [self.data]
         if not self.extended:
             args.append(self.srid)
-        self.function_expr = functions.Function(*args)
+        self.function_expr = getattr(functions.func, func_name)(*args)
 
     @staticmethod
     def _data_from_desc(desc):

diff --git a/tests/test_functional.py b/tests/test_functional.py
@@ -316,12 +316,14 @@ def setup(self):
 
     def teardown(self):
         session.rollback()
+        session.expunge_all()
         metadata.drop_all()
 
     def _create_one_lake(self):
         lake = Lake(WKTElement('LINESTRING(0 0,1 1)', srid=4326))
         session.add(lake)
         session.flush()
+        session.expunge_all()
         return lake.id
 
     def test_pickle_unpickle(self):
@@ -349,18 +351,21 @@ def setup(self):
 
     def teardown(self):
         session.rollback()
+        session.expunge_all()
         metadata.drop_all()
 
     def _create_one_lake(self):
         lake = Lake(WKTElement('LINESTRING(0 0,1 1)', srid=4326))
         session.add(lake)
         session.flush()
+        session.expunge_all()
         return lake.id
 
     def _create_one_poi(self):
         p = Poi('POINT(5 45)')
         session.add(p)
         session.flush()
+        session.expunge_all()
         return p.id
 
     def test_ST_GeometryType(self):
@@ -390,6 +395,7 @@ def test_ST_Buffer(self):
         assert isinstance(r1, WKBElement)
 
         lake = session.query(Lake).get(lake_id)
+        assert isinstance(lake.geom, WKBElement)
         r2 = session.execute(lake.geom.ST_Buffer(2)).scalar()
         assert isinstance(r2, WKBElement)
 
@@ -407,6 +413,7 @@ def test_ST_Buffer(self):
     def test_ST_Dump(self):
         lake_id = self._create_one_lake()
         lake = session.query(Lake).get(lake_id)
+        assert isinstance(lake.geom, WKBElement)
 
         s = select([func.ST_Dump(Lake.__table__.c.geom)])
         r1 = session.execute(s).scalar()
@@ -439,6 +446,7 @@ def test_ST_Dump(self):
     def test_ST_DumpPoints(self):
         lake_id = self._create_one_lake()
         lake = session.query(Lake).get(lake_id)
+        assert isinstance(lake.geom, WKBElement)
 
         dump = lake.geom.ST_DumpPoints()
 
@@ -486,6 +494,7 @@ def test_comparator_case_insensitivity(self):
         assert isinstance(r1, WKBElement)
 
         lake = session.query(Lake).get(lake_id)
+        assert isinstance(lake.geom, WKBElement)
 
         r2 = session.execute(lake.geom.ST_Buffer(2)).scalar()
         assert isinstance(r2, WKBElement)