code cleaning

aequilibrae/paths/graph.py

@@ -1,14 +1,16 @@
-from os.path import join
 import pickle
 import uuid
+from abc import ABC
 from datetime import datetime
-from abc import ABC, abstractmethod
+from os.path import join
 from typing import List, Tuple, Optional
+
 import numpy as np
 import pandas as pd
-from aequilibrae.context import get_logger
 from aequilibrae.paths.AoN import build_compressed_graph
 
+from aequilibrae.context import get_logger
+
 
 class GraphBase(ABC):
     """

aequilibrae/project/basic_table.py

@@ -3,6 +3,7 @@
 
 from aequilibrae.project.field_editor import FieldEditor
 from aequilibrae.utils.db_utils import commit_and_close
+from aequilibrae.utils.spatialite_utils import connect_spatialite
 
 
 class BasicTable:
@@ -34,3 +35,6 @@ def __copy__(self):
 
     def __deepcopy__(self, memodict=None):
         raise Exception(f"{self.__table_type__} object cannot be copied")
+
+    def conn(self):
+        return commit_and_close(connect_spatialite(self.project.path_to_file))

aequilibrae/project/network/links.py

@@ -7,8 +7,6 @@
 from aequilibrae.project.data_loader import DataLoader
 from aequilibrae.project.network.link import Link
 from aequilibrae.project.table_loader import TableLoader
-from aequilibrae.utils.db_utils import commit_and_close
-from aequilibrae.utils.spatialite_utils import connect_spatialite
 
 
 class Links(BasicTable):
@@ -117,7 +115,7 @@ def delete(self, link_id: int) -> None:
             link = self.__items.pop(link_id)  # type: Link
             link.delete()
         else:
-            with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+            with self.conn as conn:
                 d = conn.execute("Delete from Links where link_id=?", [link_id]).rowcount
         if d:
             self.project.logger.warning(f"Link {link_id} was successfully removed from the project database")
@@ -127,7 +125,7 @@ def delete(self, link_id: int) -> None:
     def refresh_fields(self) -> None:
         """After adding a field one needs to refresh all the fields recognized by the software"""
         tl = TableLoader()
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             self.__max_id = conn.execute("select coalesce(max(link_id),0) from Links").fetchone()[0]
             tl.load_structure(conn, "links")
         self.sql = tl.sql
@@ -160,7 +158,7 @@ def __existence_error(self, link_id):
         raise ValueError(f"Link {link_id} does not exist in the model")
 
     def __link_data(self, link_id: int) -> dict:
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             data = conn.execute(f"{self.sql} where link_id=?", [link_id]).fetchone()
         if data:
             return {key: val for key, val in zip(self.__fields, data)}

aequilibrae/project/network/nodes.py

@@ -6,8 +6,6 @@
 from aequilibrae.project.data_loader import DataLoader
 from aequilibrae.project.network.node import Node
 from aequilibrae.project.table_loader import TableLoader
-from aequilibrae.utils.db_utils import commit_and_close
-from aequilibrae.utils.spatialite_utils import connect_spatialite
 
 
 class Nodes(BasicTable):
@@ -63,7 +61,7 @@ def get(self, node_id: int) -> Node:
             else:
                 self.__items[node.node_id] = self.__items.pop(node_id)
 
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             data = conn.execute(f"{self.sql} where node_id=?", [node_id]).fetchone()
         if data:
             data = {key: val for key, val in zip(self.__fields, data)}
@@ -76,7 +74,7 @@ def get(self, node_id: int) -> Node:
     def refresh_fields(self) -> None:
         """After adding a field one needs to refresh all the fields recognized by the software"""
         tl = TableLoader()
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             tl.load_structure(conn, "nodes")
         self.sql = tl.sql
         self.__fields = deepcopy(tl.fields)
@@ -94,7 +92,7 @@ def new_centroid(self, node_id: int) -> Node:
             **node_id** (:obj:`int`): Id of the centroid to be created
         """
 
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             ct = conn.execute("select count(*) from nodes where node_id=?", [node_id]).fetchone()[0]
         if ct > 0:
             raise Exception("Node_id already exists. Failed to create it")
@@ -127,7 +125,7 @@ def lonlat(self) -> pd.DataFrame:
         :Returns:
             **table** (:obj:`DataFrame`): Pandas DataFrame with all the nodes, with geometry as lon/lat
         """
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             df = pd.read_sql("SELECT node_id, ST_X(geometry) AS lon, ST_Y(geometry) AS lat FROM nodes", conn)
         return df
 

aequilibrae/project/network/periods.py

@@ -66,8 +66,8 @@ def get(self, period_id: int) -> Period:
             else:
                 self.__items[period.period_id] = self.__items.pop(period_id)
 
-        self._curr.execute(f"{self.sql} where period_id=?", [period_id])
-        data = self._curr.fetchone()
+        with self.conn as conn:
+            data = conn.execute(f"{self.sql} where period_id=?", [period_id]).fetchone()
         if data:
             data = {key: val for key, val in zip(self.__fields, data)}
             period = Period(data, self.project)
@@ -79,7 +79,8 @@ def get(self, period_id: int) -> Period:
     def refresh_fields(self) -> None:
         """After adding a field one needs to refresh all the fields recognized by the software"""
         tl = TableLoader()
-        tl.load_structure(self._curr, "periods")
+        with self.conn as conn:
+            tl.load_structure(conn, "periods")
         self.sql = tl.sql
         self.__fields = deepcopy(tl.fields)
 
@@ -99,8 +100,9 @@ def new_period(self, period_id: int, start: int, end: int, description: str = No
             **description** (:obj:`str`): Optional human readable description of the time period e.g. '1pm - 5pm'
         """
 
-        self._curr.execute("SELECT COUNT(*) FROM periods WHERE period_id=?", [period_id])
-        if self._curr.fetchone()[0] > 0:
+        with self.conn as conn:
+            dt = conn.execute("SELECT COUNT(*) FROM periods WHERE period_id=?", [period_id]).fetchone()[0]
+        if dt[0] > 0:
             raise Exception("period_id already exists. Failed to create it")
 
         data = {key: None for key in self.__fields}
@@ -123,7 +125,8 @@ def data(self) -> pd.DataFrame:
         :Returns:
             **table** (:obj:`DataFrame`): Pandas DataFrame with all the periods
         """
-        dl = DataLoader(self.conn, "periods")
+        with self.conn as conn:
+            dl = DataLoader(conn, "periods")
         return dl.load_table()
 
     def __del__(self):

aequilibrae/project/zoning.py

@@ -10,9 +10,7 @@
 from aequilibrae.project.project_creation import run_queries_from_sql_file
 from aequilibrae.project.table_loader import TableLoader
 from aequilibrae.project.zone import Zone
-from aequilibrae.utils.db_utils import commit_and_close
 from aequilibrae.utils.geo_index import GeoIndex
-from aequilibrae.utils.spatialite_utils import connect_spatialite
 
 
 class Zoning(BasicTable):
@@ -73,7 +71,7 @@ def create_zoning_layer(self):
 
         if not self.__has_zoning():
             qry_file = join(realpath(__file__), "database_specification", "tables", "zones.sql")
-            with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+            with self.conn as conn:
                 run_queries_from_sql_file(conn, self.project.logger, qry_file)
             self.__load()
         else:
@@ -85,7 +83,7 @@ def coverage(self) -> Polygon:
         :Returns:
             **model coverage** (:obj:`Polygon`): Shapely (Multi)polygon of the zoning system.
         """
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             dt = conn.execute('Select ST_asBinary("geometry") from zones;').fetchall()
         polygons = [shapely.wkb.loads(x[0]) for x in dt]
         return unary_union(polygons)
@@ -132,13 +130,13 @@ def refresh_geo_index(self):
             self.__geo_index.insert(feature_id=zone_id, geometry=zone.geometry)
 
     def __has_zoning(self):
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             dt = conn.execute("SELECT name FROM sqlite_master WHERE type='table';").fetchall()
         return any(["zone" in x[0].lower() for x in dt])
 
     def __load(self):
         tl = TableLoader()
-        with commit_and_close(connect_spatialite(self.project.path_to_file)) as conn:
+        with self.conn as conn:
             zones_list = tl.load_table(conn, "zones")
         self.__fields = deepcopy(tl.fields)