Merge pull request #307 from capitalone/develop

Release v0.6.2

locopy/_version.py

@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.6.1"
+__version__ = "0.6.2"

locopy/database.py

@@ -225,7 +225,7 @@ def to_dataframe(self, df_type="pandas", size=None):
             return None
 
         if df_type == "pandas":
-            return pandas.DataFrame(fetched, columns=columns)
+            return pandas.DataFrame(fetched, columns=columns or None)
         elif df_type == "polars":
             return polars.DataFrame(fetched, schema=columns, orient="row")
 

pyproject.toml

@@ -6,7 +6,7 @@ authors = [
   { name="Faisal Dosani", email="faisal.dosani@capitalone.com" },
 ]
 license = {text = "Apache Software License"}
-dependencies = ["boto3<=1.35.9,>=1.9.92", "PyYAML<=6.0.1,>=5.1", "pandas<=2.2.2,>=0.25.2", "numpy<=2.0.2,>=1.22.0", "polars>=0.20.0"]
+dependencies = ["boto3<=1.35.43,>=1.9.92", "PyYAML<=6.0.1,>=5.1", "pandas<=2.2.3,>=0.25.2", "numpy<=2.0.2,>=1.22.0", "polars>=0.20.0"]
 
 requires-python = ">=3.9.0"
 classifiers = [

tests/test_database.py

@@ -23,7 +23,10 @@
 import sqlite3
 from unittest import mock
 
+import pandas as pd
 import pg8000
+import polars as pl
+import polars.testing as pltest
 import psycopg2
 import pytest
 import snowflake.connector
@@ -220,37 +223,47 @@ def test_execute_sql_exception(credentials, dbapi):
 
 
 @pytest.mark.parametrize("dbapi", DBAPIS)
-@mock.patch("pandas.DataFrame")
-def test_to_dataframe_all(mock_pandas, credentials, dbapi):
+def test_to_dataframe_all_pandas(credentials, dbapi):
     with mock.patch(dbapi.__name__ + ".connect") as mock_connect:
         mock_connect.return_value.cursor.return_value.fetchall.return_value = [
             (1, 2),
             (2, 3),
             (3,),
         ]
+        expected_df = pd.DataFrame(
+            [
+                (1, 2),
+                (2, 3),
+                (3,),
+            ]
+        )
         with Database(dbapi=dbapi, **credentials) as test:
             test.execute("SELECT 'hello world' AS fld")
-            df = test.to_dataframe()
-
+            df = test.to_dataframe(df_type="pandas")
+        pd.testing.assert_frame_equal(df, expected_df)
     assert mock_connect.return_value.cursor.return_value.fetchall.called
-    mock_pandas.assert_called_with(test.cursor.fetchall(), columns=[])
 
 
 @pytest.mark.parametrize("dbapi", DBAPIS)
-@mock.patch("pandas.DataFrame")
-def test_to_dataframe_custom_size(mock_pandas, credentials, dbapi):
+def test_to_dataframe_custom_size(credentials, dbapi):
     with mock.patch(dbapi.__name__ + ".connect") as mock_connect:
         mock_connect.return_value.cursor.return_value.fetchmany.return_value = [
             (1, 2),
             (2, 3),
             (3,),
         ]
+        expected_df = pd.DataFrame(
+            [
+                (1, 2),
+                (2, 3),
+                (3,),
+            ]
+        )
         with Database(dbapi=dbapi, **credentials) as test:
             test.execute("SELECT 'hello world' AS fld")
             df = test.to_dataframe(size=5)
-
+        pd.testing.assert_frame_equal(df, expected_df)
     mock_connect.return_value.cursor.return_value.fetchmany.assert_called_with(5)
-    mock_pandas.assert_called_with(test.cursor.fetchmany(), columns=[])
 
 
 @pytest.mark.parametrize("dbapi", DBAPIS)
@@ -264,22 +277,21 @@ def test_to_dataframe_none(mock_pandas, credentials, dbapi):
             mock_pandas.assert_not_called()
 
 
-# TODO: remove dataframe mocking
 @pytest.mark.parametrize("dbapi", DBAPIS)
-@mock.patch("polars.DataFrame")
-def test_to_dataframe_all_polars(mock_polars, credentials, dbapi):
+def test_to_dataframe_all_polars(credentials, dbapi):
     with mock.patch(dbapi.__name__ + ".connect") as mock_connect:
         mock_connect.return_value.cursor.return_value.fetchall.return_value = [
             (1, 2),
             (2, 3),
             (3, 4),
         ]
+        expected_df = pl.DataFrame([[1, 2, 3], [2, 3, 4]])
         with Database(dbapi=dbapi, **credentials) as test:
             test.execute("SELECT 'hello world' AS fld")
             df = test.to_dataframe(df_type="polars")
+        pltest.assert_frame_equal(df, expected_df)
 
     assert mock_connect.return_value.cursor.return_value.fetchall.called
-    mock_polars.assert_called_with(test.cursor.fetchall(), schema=[], orient="row")
 
 
 @pytest.mark.parametrize("dbapi", DBAPIS)