fixed lint issues

Author: ryanseq-gyg

dataframe_expectations/sanity_checks.py

@@ -104,9 +104,7 @@ def _extract_expectation_name(self, decorator) -> Optional[str]:
         if isinstance(decorator, ast.Call) and decorator.args:
             first_arg = decorator.args[0]
             if isinstance(first_arg, ast.Constant):
-                return first_arg.value
-            elif isinstance(first_arg, ast.Str):  # Python < 3.8 compatibility
-                return first_arg.s
+                return str(first_arg.value)
         return None
 
     def _discover_suite_methods(self):

tests/expectations_helper_classes/test_column_expectations.py

@@ -25,9 +25,9 @@ def test_get_expectation_name(expectation):
     """
     Test that the expectation name is the class name.
     """
-    assert (
-        expectation.get_expectation_name() == "MyColumnExpectation"
-    ), f"Expected 'MyColumnExpectation' but got: {expectation.get_expectation_name()}"
+    assert expectation.get_expectation_name() == "MyColumnExpectation", (
+        f"Expected 'MyColumnExpectation' but got: {expectation.get_expectation_name()}"
+    )
 
 
 def test_validate_for_pandas_df(expectation):

tests/expectations_helper_classes/test_expectation_registry.py

@@ -33,9 +33,9 @@ def dummy_expectation_factory(foo=None):
         return DummyExpectation(foo=foo)
 
     instance = DataFrameExpectationRegistry.get_expectation("DummyExpectation", foo=123)
-    assert isinstance(
-        instance, DummyExpectation
-    ), f"Expected DummyExpectation instance but got: {type(instance)}"
+    assert isinstance(instance, DummyExpectation), (
+        f"Expected DummyExpectation instance but got: {type(instance)}"
+    )
     assert instance.foo == 123, f"Expected foo=123 but got: {instance.foo}"
 
 
@@ -54,9 +54,9 @@ def dummy1(foo=None):
         def dummy2(foo=None):
             return DummyExpectation(foo=foo)
 
-    assert "already registered" in str(
-        context.value
-    ), f"Expected 'already registered' in error message but got: {str(context.value)}"
+    assert "already registered" in str(context.value), (
+        f"Expected 'already registered' in error message but got: {str(context.value)}"
+    )
 
 
 def test_get_unknown_expectation_raises():
@@ -65,9 +65,9 @@ def test_get_unknown_expectation_raises():
     """
     with pytest.raises(ValueError) as context:
         DataFrameExpectationRegistry.get_expectation("NonExistent")
-    assert "Unknown expectation" in str(
-        context.value
-    ), f"Expected 'Unknown expectation' in error message but got: {str(context.value)}"
+    assert "Unknown expectation" in str(context.value), (
+        f"Expected 'Unknown expectation' in error message but got: {str(context.value)}"
+    )
 
 
 def test_list_expectations():
@@ -98,16 +98,16 @@ def dummy(foo=None):
         return DummyExpectation(foo=foo)
 
     names_before = DataFrameExpectationRegistry.list_expectations()
-    assert (
-        "ToRemove" in names_before
-    ), f"Expected 'ToRemove' in expectations list before removal but got: {names_before}"
+    assert "ToRemove" in names_before, (
+        f"Expected 'ToRemove' in expectations list before removal but got: {names_before}"
+    )
 
     DataFrameExpectationRegistry.remove_expectation("ToRemove")
 
     names_after = DataFrameExpectationRegistry.list_expectations()
-    assert (
-        "ToRemove" not in names_after
-    ), f"Expected 'ToRemove' not in expectations list after removal but got: {names_after}"
+    assert "ToRemove" not in names_after, (
+        f"Expected 'ToRemove' not in expectations list after removal but got: {names_after}"
+    )
 
 
 def test_remove_nonexistent_expectation_raises():
@@ -116,6 +116,6 @@ def test_remove_nonexistent_expectation_raises():
     """
     with pytest.raises(ValueError) as context:
         DataFrameExpectationRegistry.remove_expectation("DefinitelyNotThere")
-    assert "not found" in str(
-        context.value
-    ), f"Expected 'not found' in error message but got: {str(context.value)}"
+    assert "not found" in str(context.value), (
+        f"Expected 'not found' in error message but got: {str(context.value)}"
+    )

tests/expectations_helper_classes/test_expectations.py

@@ -35,28 +35,28 @@ def test_data_frame_type_enum():
     """
     Test that the DataFrameType enum has the correct values.
     """
-    assert (
-        DataFrameType.PANDAS.value == "pandas"
-    ), f"Expected 'pandas' but got: {DataFrameType.PANDAS.value}"
-    assert (
-        DataFrameType.PYSPARK.value == "pyspark"
-    ), f"Expected 'pyspark' but got: {DataFrameType.PYSPARK.value}"
+    assert DataFrameType.PANDAS.value == "pandas", (
+        f"Expected 'pandas' but got: {DataFrameType.PANDAS.value}"
+    )
+    assert DataFrameType.PYSPARK.value == "pyspark", (
+        f"Expected 'pyspark' but got: {DataFrameType.PYSPARK.value}"
+    )
 
     # Test string comparison (now works directly!)
     assert DataFrameType.PANDAS == "pandas", "Expected DataFrameType.PANDAS == 'pandas' to be True"
-    assert (
-        DataFrameType.PYSPARK == "pyspark"
-    ), "Expected DataFrameType.PYSPARK == 'pyspark' to be True"
+    assert DataFrameType.PYSPARK == "pyspark", (
+        "Expected DataFrameType.PYSPARK == 'pyspark' to be True"
+    )
 
 
 def test_get_expectation_name():
     """
     Test that the expectation name is the class name.
     """
     expectation = MyTestExpectation()
-    assert (
-        expectation.get_expectation_name() == "MyTestExpectation"
-    ), f"Expected 'MyTestExpectation' but got: {expectation.get_expectation_name()}"
+    assert expectation.get_expectation_name() == "MyTestExpectation", (
+        f"Expected 'MyTestExpectation' but got: {expectation.get_expectation_name()}"
+    )
 
 
 def test_validate_unsupported_dataframe_type():
@@ -150,55 +150,55 @@ def test_infer_data_frame_type(spark):
     # Test pandas DataFrame
     pandas_df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
     data_frame_type = expectation.infer_data_frame_type(pandas_df)
-    assert (
-        data_frame_type == DataFrameType.PANDAS
-    ), f"Expected PANDAS type but got: {data_frame_type}"
+    assert data_frame_type == DataFrameType.PANDAS, (
+        f"Expected PANDAS type but got: {data_frame_type}"
+    )
 
     # Test PySpark DataFrame
     spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["col1", "col2"])
     data_frame_type = expectation.infer_data_frame_type(spark_df)
-    assert (
-        data_frame_type == DataFrameType.PYSPARK
-    ), f"Expected PYSPARK type but got: {data_frame_type}"
+    assert data_frame_type == DataFrameType.PYSPARK, (
+        f"Expected PYSPARK type but got: {data_frame_type}"
+    )
 
     # Test empty pandas DataFrame
     empty_pandas_df = pd.DataFrame(columns=["col1", "col2"])
     data_frame_type = expectation.infer_data_frame_type(empty_pandas_df)
-    assert (
-        data_frame_type == DataFrameType.PANDAS
-    ), f"Expected PANDAS type for empty DataFrame but got: {data_frame_type}"
+    assert data_frame_type == DataFrameType.PANDAS, (
+        f"Expected PANDAS type for empty DataFrame but got: {data_frame_type}"
+    )
 
     # Test empty PySpark DataFrame
     empty_spark_df = spark.createDataFrame([], "col1 INT, col2 STRING")
     data_frame_type = expectation.infer_data_frame_type(empty_spark_df)
-    assert (
-        data_frame_type == DataFrameType.PYSPARK
-    ), f"Expected PYSPARK type for empty DataFrame but got: {data_frame_type}"
+    assert data_frame_type == DataFrameType.PYSPARK, (
+        f"Expected PYSPARK type for empty DataFrame but got: {data_frame_type}"
+    )
 
     # Test unsupported DataFrame types
     with pytest.raises(ValueError) as context:
         expectation.infer_data_frame_type(None)
-    assert "Unsupported DataFrame type" in str(
-        context.value
-    ), f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    assert "Unsupported DataFrame type" in str(context.value), (
+        f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    )
 
     with pytest.raises(ValueError) as context:
         expectation.infer_data_frame_type("not_a_dataframe")
-    assert "Unsupported DataFrame type" in str(
-        context.value
-    ), f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    assert "Unsupported DataFrame type" in str(context.value), (
+        f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    )
 
     with pytest.raises(ValueError) as context:
         expectation.infer_data_frame_type([1, 2, 3])
-    assert "Unsupported DataFrame type" in str(
-        context.value
-    ), f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    assert "Unsupported DataFrame type" in str(context.value), (
+        f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    )
 
     with pytest.raises(ValueError) as context:
         expectation.infer_data_frame_type({"col1": [1, 2, 3]})
-    assert "Unsupported DataFrame type" in str(
-        context.value
-    ), f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    assert "Unsupported DataFrame type" in str(context.value), (
+        f"Expected 'Unsupported DataFrame type' in error message but got: {str(context.value)}"
+    )
 
     # Test with objects that might have similar attributes but aren't DataFrames
     class FakeDataFrame:
@@ -237,9 +237,9 @@ def test_infer_data_frame_type_with_connect_dataframe_available():
 
         # Test that Connect DataFrame is identified as PYSPARK type
         data_frame_type = expectation.infer_data_frame_type(mock_connect_df)
-        assert (
-            data_frame_type == DataFrameType.PYSPARK
-        ), f"Expected PYSPARK type for Connect DataFrame but got: {data_frame_type}"
+        assert data_frame_type == DataFrameType.PYSPARK, (
+            f"Expected PYSPARK type for Connect DataFrame but got: {data_frame_type}"
+        )
 
 
 @patch("dataframe_expectations.expectations.PySparkConnectDataFrame", None)
@@ -252,15 +252,15 @@ def test_infer_data_frame_type_without_connect_support(spark):
     # Test that regular DataFrames still work when Connect is not available
     pandas_df = pd.DataFrame({"col1": [1, 2, 3]})
     data_frame_type = expectation.infer_data_frame_type(pandas_df)
-    assert (
-        data_frame_type == DataFrameType.PANDAS
-    ), f"Expected PANDAS type but got: {data_frame_type}"
+    assert data_frame_type == DataFrameType.PANDAS, (
+        f"Expected PANDAS type but got: {data_frame_type}"
+    )
 
     spark_df = spark.createDataFrame([(1,), (2,), (3,)], ["col1"])
     data_frame_type = expectation.infer_data_frame_type(spark_df)
-    assert (
-        data_frame_type == DataFrameType.PYSPARK
-    ), f"Expected PYSPARK type but got: {data_frame_type}"
+    assert data_frame_type == DataFrameType.PYSPARK, (
+        f"Expected PYSPARK type but got: {data_frame_type}"
+    )
 
 
 def test_infer_data_frame_type_connect_import_behavior(spark):
@@ -297,6 +297,6 @@ def test_infer_data_frame_type_connect_import_behavior(spark):
         # Mock Connect DataFrame should be identified as PYSPARK
         mock_connect_df = MockConnectDataFrame()
         result_type = expectation.infer_data_frame_type(mock_connect_df)
-        assert (
-            result_type == DataFrameType.PYSPARK
-        ), f"Expected PYSPARK type for Connect DataFrame but got: {result_type}"
+        assert result_type == DataFrameType.PYSPARK, (
+            f"Expected PYSPARK type for Connect DataFrame but got: {result_type}"
+        )

tests/expectations_helper_classes/test_utils.py

@@ -28,9 +28,9 @@ def func(**kwargs):
 
     with pytest.raises(ValueError) as context:
         func(a=1)
-    assert "missing required parameters" in str(
-        context.value
-    ), f"Expected 'missing required parameters' in error message but got: {str(context.value)}"
+    assert "missing required parameters" in str(context.value), (
+        f"Expected 'missing required parameters' in error message but got: {str(context.value)}"
+    )
 
 
 def test_requires_params_type_success():
@@ -56,9 +56,9 @@ def func(**kwargs):
 
     with pytest.raises(TypeError) as context:
         func(a="not-an-int", b="hello")
-    assert "type validation errors" in str(
-        context.value
-    ), f"Expected 'type validation errors' in error message but got: {str(context.value)}"
+    assert "type validation errors" in str(context.value), (
+        f"Expected 'type validation errors' in error message but got: {str(context.value)}"
+    )
 
 
 def test_requires_params_union_type_success():