Assert type of method arguments

Author: EnricoMi

Diff for: python/gresearch/spark/__init__.py

@@ -108,13 +108,19 @@ def _to_map(jvm: JVMView, map: Mapping[Any, Any]) -> JavaObject:
 
 
 def backticks(*name_parts: str) -> str:
+    for np in name_parts:
+        assert isinstance(np, str), np
     return '.'.join([f'`{part}`'
                      if '.' in part and not part.startswith('`') and not part.endswith('`')
                      else part
                      for part in name_parts])
 
 
 def distinct_prefix_for(existing: List[str]) -> str:
+    assert isinstance(existing, Iterable)
+    for e in existing:
+        assert isinstance(e, str), e
+
     # count number of suffix _ for each existing column name
     length = 1
     if existing:
@@ -128,21 +134,46 @@ def handle_configured_case_sensitivity(column_name: str, case_sensitive: bool) -
     Produces a column name that considers configured case-sensitivity of column names. When case sensitivity is
     deactivated, it lower-cases the given column name and no-ops otherwise.
     """
+    assert isinstance(column_name, str), column_name
+    assert isinstance(case_sensitive, bool), case_sensitive
+
     if case_sensitive:
         return column_name
     return column_name.lower()
 
 
 def list_contains_case_sensitivity(column_names: Iterable[str], columnName: str, case_sensitive: bool) -> bool:
+    assert isinstance(column_names, Iterable), column_names
+    for cn in column_names:
+        assert isinstance(cn, str), cn
+    assert isinstance(columnName, str), columnName
+    assert isinstance(case_sensitive, bool), case_sensitive
+
     return handle_configured_case_sensitivity(columnName, case_sensitive) in [handle_configured_case_sensitivity(c, case_sensitive) for c in column_names]
 
 
 def list_filter_case_sensitivity(column_names: Iterable[str], filter: Iterable[str], case_sensitive: bool) -> List[str]:
+    assert isinstance(column_names, Iterable), column_names
+    for cn in column_names:
+        assert isinstance(cn, str), cn
+    assert isinstance(filter, Iterable), filter
+    for f in filter:
+        assert isinstance(f, str), f
+    assert isinstance(case_sensitive, bool), case_sensitive
+
     filter_set = {handle_configured_case_sensitivity(f, case_sensitive) for f in filter}
     return [c for c in column_names if handle_configured_case_sensitivity(c, case_sensitive) in filter_set]
 
 
 def list_diff_case_sensitivity(column_names: Iterable[str], other: Iterable[str], case_sensitive: bool) -> List[str]:
+    assert isinstance(column_names, Iterable), column_names
+    for cn in column_names:
+        assert isinstance(cn, str), cn
+    assert isinstance(other, Iterable), filter
+    for o in other:
+        assert isinstance(o, str), o
+    assert isinstance(case_sensitive, bool), case_sensitive
+
     other_set = {handle_configured_case_sensitivity(f, case_sensitive) for f in other}
     return [c for c in column_names if handle_configured_case_sensitivity(c, case_sensitive) not in other_set]
 
@@ -344,6 +375,9 @@ def count_null(e: "ColumnOrName") -> Column:
     """
     if isinstance(e, str):
         e = col(e)
+    if not isinstance(e, Column):
+        raise ValueError(f"Given column must be a column name (str) or column instance (Column): {type(e)}")
+
     return count(when(e.isNull(), lit(1)))
 
 
@@ -452,7 +486,11 @@ def session_or_ctx(self: DataFrame) -> Union[SparkSession, SQLContext]:
     ConnectDataFrame.session_or_ctx = session_or_ctx
 
 
-def set_description(description: str, if_not_set: bool = False):
+def set_description(description: Optional[str], if_not_set: bool = False):
+    if description is not None:
+        assert isinstance(description, str), description
+    assert isinstance(if_not_set, bool), if_not_set
+
     context = SparkContext._active_spark_context
     jvm = _get_jvm(context)
     spark_package = jvm.uk.co.gresearch.spark.__getattr__("package$").__getattr__("MODULE$")
@@ -489,6 +527,9 @@ def job_description(description: str, if_not_set: bool = False):
 
 
 def append_description(extra_description: str, separator: str = " - "):
+    assert isinstance(extra_description, str), extra_description
+    assert isinstance(separator, str), separator
+
     context = SparkContext._active_spark_context
     jvm = _get_jvm(context)
     spark_package = jvm.uk.co.gresearch.spark.__getattr__("package$").__getattr__("MODULE$")
@@ -546,6 +587,9 @@ def install_pip_package(spark: Union[SparkSession, SparkContext], *package_or_pi
     if __version__.startswith('2.') or __version__.startswith('3.0.'):
         raise NotImplementedError(f'Not supported for PySpark __version__')
 
+    for option in package_or_pip_option:
+        assert isinstance(option, str), option
+
     # just here to assert JVM is accessible
     _get_jvm(spark)
 
@@ -591,6 +635,15 @@ def install_poetry_project(spark: Union[SparkSession, SparkContext],
     if __version__.startswith('2.') or __version__.startswith('3.0.'):
         raise NotImplementedError(f'Not supported for PySpark __version__')
 
+    for p in project:
+        assert isinstance(p, str), p
+
+    if poetry_python is not None:
+        assert isinstance(poetry_python, str), poetry_python
+    if pip_args is not None:
+        for pa in pip_args:
+            assert isinstance(pa, str), pa
+
     # just here to assert JVM is accessible
     _get_jvm(spark)
 

Diff for: python/gresearch/spark/diff/__init__.py

@@ -96,6 +96,7 @@ def with_diff_column(self, diff_column: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(diff_column, str), diff_column
         return dataclasses.replace(self, diff_column=diff_column)
 
     def with_left_column_prefix(self, left_column_prefix: str) -> 'DiffOptions':
@@ -108,6 +109,7 @@ def with_left_column_prefix(self, left_column_prefix: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(left_column_prefix, str), left_column_prefix
         return dataclasses.replace(self, left_column_prefix=left_column_prefix)
 
     def with_right_column_prefix(self, right_column_prefix: str) -> 'DiffOptions':
@@ -120,6 +122,7 @@ def with_right_column_prefix(self, right_column_prefix: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(right_column_prefix, str), right_column_prefix
         return dataclasses.replace(self, right_column_prefix=right_column_prefix)
 
     def with_insert_diff_value(self, insert_diff_value: str) -> 'DiffOptions':
@@ -132,6 +135,7 @@ def with_insert_diff_value(self, insert_diff_value: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(insert_diff_value, str), insert_diff_value
         return dataclasses.replace(self, insert_diff_value=insert_diff_value)
 
     def with_change_diff_value(self, change_diff_value: str) -> 'DiffOptions':
@@ -144,6 +148,7 @@ def with_change_diff_value(self, change_diff_value: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(change_diff_value, str), change_diff_value
         return dataclasses.replace(self, change_diff_value=change_diff_value)
 
     def with_delete_diff_value(self, delete_diff_value: str) -> 'DiffOptions':
@@ -156,6 +161,7 @@ def with_delete_diff_value(self, delete_diff_value: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(delete_diff_value, str), delete_diff_value
         return dataclasses.replace(self, delete_diff_value=delete_diff_value)
 
     def with_nochange_diff_value(self, nochange_diff_value: str) -> 'DiffOptions':
@@ -168,6 +174,7 @@ def with_nochange_diff_value(self, nochange_diff_value: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(nochange_diff_value, str), nochange_diff_value
         return dataclasses.replace(self, nochange_diff_value=nochange_diff_value)
 
     def with_change_column(self, change_column: str) -> 'DiffOptions':
@@ -180,6 +187,7 @@ def with_change_column(self, change_column: str) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(change_column, str), change_column
         return dataclasses.replace(self, change_column=change_column)
 
     def without_change_column(self) -> 'DiffOptions':
@@ -202,6 +210,7 @@ def with_diff_mode(self, diff_mode: DiffMode) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(diff_mode, DiffMode), diff_mode
         return dataclasses.replace(self, diff_mode=diff_mode)
 
     def with_sparse_mode(self, sparse_mode: bool) -> 'DiffOptions':
@@ -214,12 +223,18 @@ def with_sparse_mode(self, sparse_mode: bool) -> 'DiffOptions':
         :return: new immutable DiffOptions instance
         :rtype: DiffOptions
         """
+        assert isinstance(sparse_mode, bool), sparse_mode
         return dataclasses.replace(self, sparse_mode=sparse_mode)
 
     def with_default_comparator(self, comparator: DiffComparator) -> 'DiffOptions':
+        assert isinstance(comparator, DiffComparator), comparator
         return dataclasses.replace(self, default_comparator=comparator)
 
     def with_data_type_comparator(self, comparator: DiffComparator, *data_type: DataType) -> 'DiffOptions':
+        assert isinstance(comparator, DiffComparator), comparator
+        for dt in data_type:
+            assert isinstance(dt, DataType), dt
+
         existing_data_types = {dt.simpleString() for dt in data_type if dt in self.data_type_comparators.keys()}
         if existing_data_types:
             existing_data_types = sorted(list(existing_data_types))
@@ -231,6 +246,10 @@ def with_data_type_comparator(self, comparator: DiffComparator, *data_type: Data
         return dataclasses.replace(self, data_type_comparators=data_type_comparators)
 
     def with_column_name_comparator(self, comparator: DiffComparator, *column_name: str) -> 'DiffOptions':
+        assert isinstance(comparator, DiffComparator), comparator
+        for cn in column_name:
+            assert isinstance(cn, str), cn
+
         existing_column_names = {cn for cn in column_name if cn in self.column_name_comparators.keys()}
         if existing_column_names:
             existing_column_names = sorted(list(existing_column_names))
@@ -242,6 +261,7 @@ def with_column_name_comparator(self, comparator: DiffComparator, *column_name:
         return dataclasses.replace(self, column_name_comparators=column_name_comparators)
 
     def comparator_for(self, column: StructField) -> DiffComparator:
+        assert isinstance(column, StructField), column
         cmp = self.column_name_comparators.get(column.name)
         if cmp is None:
             cmp = self.data_type_comparators.get(column.dataType)
@@ -328,14 +348,24 @@ def diff(self, left: DataFrame, right: DataFrame, *id_or_ignore_columns: Union[s
         :type right: DataFrame
         :param id_or_ignore_columns: either id column names or two lists of column names,
                first the id column names, second the ignore column names
-        :type id_or_ignore_columns: str
+        :type *id_or_ignore_columns: str | Iterable[str]
         :return: the diff DataFrame
         :rtype DataFrame
         """
-        if len(id_or_ignore_columns) == 2 and all([isinstance(lst, Iterable) and not isinstance(lst, str) for lst in id_or_ignore_columns]):
+        assert isinstance(left, DataFrame), left
+        assert isinstance(right, DataFrame), right
+        assert isinstance(id_or_ignore_columns, (str, Iterable)), id_or_ignore_columns
+
+        if len(id_or_ignore_columns) == 2 and all(isinstance(lst, Iterable) and not isinstance(lst, str) for lst in id_or_ignore_columns):
             id_columns, ignore_columns = id_or_ignore_columns
-        else:
+            if any(not isinstance(id, str) for id in id_columns):
+                raise ValueError(f"The id_columns must all be strings: {', '.join(type(id).__name__ for id in id_columns)}")
+            if any(not isinstance(ignore, str) for ignore in ignore_columns):
+                raise ValueError(f"The ignore_columns must all be strings: {', '.join(type(ignore).__name__ for ignore in ignore_columns)}")
+        elif all(isinstance(lst, str) for lst in id_or_ignore_columns):
             id_columns, ignore_columns = (id_or_ignore_columns, [])
+        else:
+            raise ValueError(f"The id_or_ignore_columns argument must either all be strings or exactly two iterables of strings: {', '.join(type(e).__name__ for e in id_or_ignore_columns)}")
 
         return self._do_diff(left, right, id_columns, ignore_columns)
 

Diff for: python/gresearch/spark/diff/comparator/__init__.py

@@ -40,23 +40,31 @@ def nullSafeEqual() -> 'NullSafeEqualDiffComparator':
 
     @staticmethod
     def epsilon(epsilon: float) -> 'EpsilonDiffComparator':
+        assert isinstance(epsilon, float), epsilon
         return EpsilonDiffComparator(epsilon)
 
     @staticmethod
     def string(whitespace_agnostic: bool = True) -> 'StringDiffComparator':
+        assert isinstance(whitespace_agnostic, bool), whitespace_agnostic
         return StringDiffComparator(whitespace_agnostic)
 
     @staticmethod
     def duration(duration: str) -> 'DurationDiffComparator':
+        assert isinstance(duration, str), duration
         return DurationDiffComparator(duration)
 
     @staticmethod
     def map(key_type: DataType, value_type: DataType, key_order_sensitive: bool = False) -> 'MapDiffComparator':
+        assert isinstance(key_type, DataType), key_type
+        assert isinstance(value_type, DataType), value_type
+        assert isinstance(key_order_sensitive, bool), key_order_sensitive
         return MapDiffComparator(key_type, value_type, key_order_sensitive)
 
 
 class NullSafeEqualDiffComparator(DiffComparator):
     def equiv(self, left: Column, right: Column) -> Column:
+        assert isinstance(left, Column), left
+        assert isinstance(right, Column), right
         return left.eqNullSafe(right)
 
 
@@ -85,6 +93,9 @@ def as_exclusive(self) -> 'EpsilonDiffComparator':
         return dataclasses.replace(self, inclusive=False)
 
     def equiv(self, left: Column, right: Column) -> Column:
+        assert isinstance(left, Column), left
+        assert isinstance(right, Column), right
+
         threshold = greatest(abs(left), abs(right)) * self.epsilon if self.relative else lit(self.epsilon)
 
         def inclusive_epsilon(diff: Column) -> Column:
@@ -102,6 +113,8 @@ class StringDiffComparator(DiffComparator):
     whitespace_agnostic: bool
 
     def equiv(self, left: Column, right: Column) -> Column:
+        assert isinstance(left, Column), left
+        assert isinstance(right, Column), right
         return left.eqNullSafe(right)
 
 
@@ -117,6 +130,8 @@ def as_exclusive(self) -> 'DurationDiffComparator':
         return dataclasses.replace(self, inclusive=False)
 
     def equiv(self, left: Column, right: Column) -> Column:
+        assert isinstance(left, Column), left
+        assert isinstance(right, Column), right
         return left.eqNullSafe(right)
 
 
@@ -127,4 +142,6 @@ class MapDiffComparator(DiffComparator):
     key_order_sensitive: bool
 
     def equiv(self, left: Column, right: Column) -> Column:
+        assert isinstance(left, Column), left
+        assert isinstance(right, Column), right
         return left.eqNullSafe(right)