Merge pull request #3 from cid-harvard/feature/py2.7-compatibility

Python 2.7 compatibility changes

Author: makmanalp

pandas_to_postgres/_base_copy.py

@@ -1,11 +1,6 @@
 from .utilities import logger
-from io import StringIO
-from pandas import DataFrame
-from typing import Callable, List
 from sqlalchemy.schema import AddConstraint, DropConstraint
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.sql.schema import Table
-from sqlalchemy.engine.base import Connection
 
 
 class BaseCopy(object):
@@ -15,21 +10,26 @@ class BaseCopy(object):
 
     def __init__(
         self,
-        defer_sql_objs: bool = False,
-        conn: Connection = None,
-        table_obj: Table = None,
-        sql_table: str = None,
-        csv_chunksize: int = 10 ** 6,
+        defer_sql_objs=False,
+        conn=None,
+        table_obj=None,
+        sql_table=None,
+        csv_chunksize=10 ** 6,
     ):
         """
         Parameters
         ----------
-        defer_sql_objs: multiprocessing has issue with passing SQLALchemy objects, so if
+        defer_sql_objs: bool
+            multiprocessing has issue with passing SQLALchemy objects, so if
             True, defer attributing these to the object until after pickled by Pool
-        conn: SQLAlchemy connection managed outside of the object
-        table_obj: SQLAlchemy object for the destination SQL Table
-        sql_table: string of SQL table name
-        csv_chunksize: max rows to keep in memory when generating CSV for COPY
+        conn: SQLAlchemy Connection
+            Managed outside of the object
+        table_obj: SQLAlchemy Table
+            Model object for the destination SQL Table
+        sql_table: string
+            SQL table name
+        csv_chunksize: int
+            Max rows to keep in memory when generating CSV for COPY
         """
 
         self.rows = 0
@@ -47,8 +47,10 @@ def instantiate_sql_objs(self, conn, table_obj):
 
         Parameters
         ----------
-        conn: SQLAlchemy connection managed outside of the object
-        table_obj: SQLAlchemy object for the destination SQL Table
+        conn: SQLAlchemy Connection
+            Managed outside of the object
+        table_obj: SQLAlchemy Table
+            Model object for the destination SQL Table
         """
         self.conn = conn
         self.table_obj = table_obj
@@ -61,71 +63,77 @@ def drop_pk(self):
         Drop primary key constraints on PostgreSQL table as well as CASCADE any other
         constraints that may rely on the PK
         """
-        logger.info(f"Dropping {self.sql_table} primary key")
+        logger.info("Dropping {} primary key".format(self.sql_table))
         try:
             with self.conn.begin_nested():
                 self.conn.execute(DropConstraint(self.primary_key, cascade=True))
         except SQLAlchemyError:
-            logger.info(f"{self.sql_table} primary key not found. Skipping")
+            logger.info("{} primary key not found. Skipping".format(self.sql_table))
 
     def create_pk(self):
         """Create primary key constraints on PostgreSQL table"""
-        logger.info(f"Creating {self.sql_table} primary key")
+        logger.info("Creating {} primary key".format(self.sql_table))
         self.conn.execute(AddConstraint(self.primary_key))
 
     def drop_fks(self):
         """Drop foreign key constraints on PostgreSQL table"""
         for fk in self.foreign_keys:
-            logger.info(f"Dropping foreign key {fk.name}")
+            logger.info("Dropping foreign key {}".format(fk.name))
             try:
                 with self.conn.begin_nested():
                     self.conn.execute(DropConstraint(fk))
             except SQLAlchemyError:
-                logger.warn(f"Foreign key {fk.name} not found")
+                logger.warn("Foreign key {} not found".format(fk.name))
 
     def create_fks(self):
         """Create foreign key constraints on PostgreSQL table"""
         for fk in self.foreign_keys:
             try:
-                logger.info(f"Creating foreign key {fk.name}")
+                logger.info("Creating foreign key {fk.name}".format(fk.name))
                 self.conn.execute(AddConstraint(fk))
             except SQLAlchemyError:
-                logger.warn(f"Error creating foreign key {fk.name}")
+                logger.warn("Error creating foreign key {fk.name}".format(fk.name))
 
     def truncate(self):
         """TRUNCATE PostgreSQL table"""
-        logger.info(f"Truncating {self.sql_table}")
-        self.conn.execute(f"TRUNCATE TABLE {self.sql_table};")
+        logger.info("Truncating {}".format(self.sql_table))
+        self.conn.execute("TRUNCATE TABLE {};".format(self.sql_table))
 
     def analyze(self):
         """Run ANALYZE on PostgreSQL table"""
-        logger.info(f"Analyzing {self.sql_table}")
-        self.conn.execute(f"ANALYZE {self.sql_table};")
+        logger.info("Analyzing {}".format(self.sql_table))
+        self.conn.execute("ANALYZE {};".format(self.sql_table))
 
-    def copy_from_file(self, file_object: StringIO):
+    def copy_from_file(self, file_object):
         """
         COPY to PostgreSQL table using StringIO CSV object
 
         Parameters
         ----------
-        file_object: CSV formatted data to COPY from DataFrame to PostgreSQL
+        file_object: StringIO
+            CSV formatted data to COPY from DataFrame to PostgreSQL
         """
         cur = self.conn.connection.cursor()
         file_object.seek(0)
         columns = file_object.readline()
-        sql = f"COPY {self.sql_table} ({columns}) FROM STDIN WITH CSV FREEZE"
+        sql = "COPY {table} ({columns}) FROM STDIN WITH CSV FREEZE".format(
+            table=self.sql_table, columns=columns
+        )
         cur.copy_expert(sql=sql, file=file_object)
 
-    def data_formatting(self, df: DataFrame, functions: List[Callable] = [], **kwargs):
+    def data_formatting(self, df, functions=[], **kwargs):
         """
         Call each function in the functions list arg on the DataFrame and return
 
         Parameters
         ----------
-        df: dataframe to format
-        functions: list of functions to apply to df. each gets passed df, self as
-            copy_obj, and all kwargs passed to data_formatting
-        **kwargs: kwargs to pass on to each function
+        df: pandas DataFrame
+            dataframe to format
+        functions: list of functions
+            Functions to apply to df. each gets passed df, self as copy_obj, and all
+            kwargs passed to data_formatting
+        **kwargs
+            kwargs to pass on to each function
         """
         for f in functions:
             df = f(df, copy_obj=self, **kwargs)

pandas_to_postgres/copy_df.py

@@ -1,20 +1,31 @@
 from .utilities import create_file_object, df_generator, logger, cast_pandas
 from ._base_copy import BaseCopy
 
-import pandas as pd
-from sqlalchemy.sql.schema import Table
-from sqlalchemy.engine.base import Connection
-
 
 class DataFrameCopy(BaseCopy):
+    """
+    Class for handling a standard case of iterating over a pandas DataFrame in chunks
+    and COPYing to PostgreSQL via StringIO CSV
+    """
+
     def __init__(
-        self,
-        df: pd.DataFrame,
-        defer_sql_objs: bool = False,
-        conn: Connection = None,
-        table_obj: Table = None,
-        csv_chunksize: int = 10 ** 6,
+        self, df, defer_sql_objs=False, conn=None, table_obj=None, csv_chunksize=10 ** 6
     ):
+        """
+        Parameters
+        ----------
+        df: pandas DataFrame
+            Data to copy to database table
+        defer_sql_objs: bool
+            multiprocessing has issue with passing SQLALchemy objects, so if
+            True, defer attributing these to the object until after pickled by Pool
+        conn: SQlAlchemy Connection
+            Managed outside of the object
+        table_obj: SQLAlchemy model object
+            Destination SQL Table
+        csv_chunksize: int
+            Max rows to keep in memory when generating CSV for COPY
+        """
         super().__init__(defer_sql_objs, conn, table_obj, csv_chunksize)
 
         self.df = df
@@ -37,7 +48,7 @@ def copy(self, functions=[cast_pandas]):
                 self.copy_from_file(fo)
                 del fo
 
-            logger.info(f"All chunks copied ({self.rows} rows)")
+            logger.info("All chunks copied ({} rows)".format(self.rows))
 
         self.create_pk()
         self.create_fks()