diff --git a/quinn/schema_helpers.py b/quinn/schema_helpers.py index 8fd18b1a..d314ff1b 100644 --- a/quinn/schema_helpers.py +++ b/quinn/schema_helpers.py @@ -132,7 +132,7 @@ def _lookup_type(type_str: str) -> T.DataType: return type_lookup[type_str] - def _convert_nullable(null_str: str) -> bool: + def _convert_nullable(null_str: Optional[str]) -> bool: if null_str is None: return True @@ -159,8 +159,8 @@ def _convert_nullable(null_str: str) -> bool: field = T.StructField( name=row["name"], dataType=_lookup_type(row["type"]), - nullable=_convert_nullable(row.get("nullable", True)), - metadata=_validate_json(row.get("metadata", None)), + nullable=_convert_nullable(row["nullable"]) if "nullable" in row else True, + metadata=_validate_json(row["metadata"] if "metadata" in row else None), # noqa: SIM401 ) fields.append(field) @@ -175,4 +175,8 @@ def complex_fields(schema: T.StructType) -> dict[str, object]: :return: A dictionary with complex field names as keys and their respective data types as values. :rtype: Dict[str, object] """ - return {field.name: field.dataType for field in schema.fields if isinstance(field.dataType, (T.ArrayType, T.StructType, T.MapType))} + return { + field.name: field.dataType + for field in schema.fields + if isinstance(field.dataType, (T.ArrayType, T.StructType, T.MapType)) + }