only allow enums values

Author: Linchin

google/cloud/bigquery/enums.py

@@ -482,7 +482,7 @@ class SourceColumnMatch(str, enum.Enum):
     columns to match the field names in the schema."""
 
 
-class TimestampPrecision(object):
+class TimestampPrecision(enum.Enum):
     """Precision (maximum number of total digits in base 10) for seconds of
     TIMESTAMP type."""
 

google/cloud/bigquery/schema.py

@@ -221,7 +221,7 @@ def __init__(
         range_element_type: Union[FieldElementType, str, None] = None,
         rounding_mode: Union[enums.RoundingMode, str, None] = None,
         foreign_type_definition: Optional[str] = None,
-        timestamp_precision: Union[enums.TimestampPrecision, int, None] = None,
+        timestamp_precision: Optional[enums.TimestampPrecision] = None,
     ):
         self._properties: Dict[str, Any] = {
             "name": name,
@@ -246,8 +246,13 @@ def __init__(
                 if isinstance(policy_tags, PolicyTagList)
                 else None
             )
-        if timestamp_precision is not None:
-            self._properties["timestampPrecision"] = timestamp_precision
+        if isinstance(timestamp_precision, enums.TimestampPrecision):
+            self._properties["timestampPrecision"] = timestamp_precision.value
+        elif timestamp_precision is not None:
+            raise ValueError(
+                "timestamp_precision must be class enums.TimestampPrecision "
+                f"or None, got {type(timestamp_precision)} instead."
+            )
         if isinstance(range_element_type, str):
             self._properties["rangeElementType"] = {"type": range_element_type}
         if isinstance(range_element_type, FieldElementType):
@@ -390,7 +395,7 @@ def timestamp_precision(self):
         """Union[enums.TimestampPrecision, int, None]: Precision (maximum number
         of total digits in base 10) for seconds of TIMESTAMP type.
         """
-        return _helpers._int_or_none(self._properties.get("timestampPrecision"))
+        return enums.TimestampPrecision(self._properties.get("timestampPrecision"))
 
     def to_api_repr(self) -> dict:
         """Return a dictionary representing this schema field.
@@ -426,6 +431,8 @@ def _key(self):
             None if self.policy_tags is None else tuple(sorted(self.policy_tags.names))
         )
 
+        timestamp_precision = self._properties.get("timestampPrecision")
+
         return (
             self.name,
             field_type,
@@ -435,7 +442,7 @@ def _key(self):
             self.description,
             self.fields,
             policy_tags,
-            self.timestamp_precision,
+            timestamp_precision,
         )
 
     def to_standard_sql(self) -> standard_sql.StandardSqlField:
@@ -548,9 +555,11 @@ def _to_schema_fields(schema):
     if isinstance(schema, Sequence):
         # Input is a Sequence (e.g. a list): Process and return a list of SchemaFields
         return [
-            field
-            if isinstance(field, SchemaField)
-            else SchemaField.from_api_repr(field)
+            (
+                field
+                if isinstance(field, SchemaField)
+                else SchemaField.from_api_repr(field)
+            )
             for field in schema
         ]
 

tests/unit/test_schema.py

@@ -52,6 +52,9 @@ def test_constructor_defaults(self):
         self.assertIsNone(field.default_value_expression)
         self.assertEqual(field.rounding_mode, None)
         self.assertEqual(field.foreign_type_definition, None)
+        self.assertEqual(
+            field.timestamp_precision, enums.TimestampPrecision.MICROSECOND
+        )
 
     def test_constructor_explicit(self):
         FIELD_DEFAULT_VALUE_EXPRESSION = "This is the default value for this field"
@@ -89,7 +92,7 @@ def test_constructor_explicit(self):
         self.assertEqual(field.rounding_mode, "ROUNDING_MODE_UNSPECIFIED")
         self.assertEqual(field.foreign_type_definition, "INTEGER")
         self.assertEqual(
-            field._properties["timestampPrecision"],
+            field.timestamp_precision,
             enums.TimestampPrecision.PICOSECOND,
         )
 
@@ -207,7 +210,7 @@ def test_to_api_repr_w_timestamp_precision(self):
                 "mode": "NULLABLE",
                 "name": "foo",
                 "type": "TIMESTAMP",
-                "timestampPrecision": enums.TimestampPrecision.PICOSECOND,
+                "timestampPrecision": 12,
             },
         )