BigQuery: Add StandardSqlDataTypes enum to BigQuery

bigquery/docs/reference.rst

@@ -164,6 +164,16 @@ Magics
 
     magics
 
+
+Enums
+=====
+
+.. autosummary::
+    :toctree: generated
+
+    enums.StandardSqlDataTypes
+
+
 Additional Types
 ================
 

bigquery/google/cloud/bigquery/__init__.py

@@ -36,6 +36,7 @@
 from google.cloud.bigquery.dataset import AccessEntry
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetReference
+from google.cloud.bigquery.enums import StandardSqlDataTypes
 from google.cloud.bigquery.external_config import ExternalConfig
 from google.cloud.bigquery.external_config import BigtableOptions
 from google.cloud.bigquery.external_config import BigtableColumnFamily
@@ -130,6 +131,7 @@
     "Encoding",
     "QueryPriority",
     "SchemaUpdateOption",
+    "StandardSqlDataTypes",
     "SourceFormat",
     "WriteDisposition",
 ]

bigquery/google/cloud/bigquery/enums.py

@@ -0,0 +1,69 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+import enum
+import six
+
+from google.cloud.bigquery_v2.gapic import enums as gapic_enums
+
+
+_SQL_SCALAR_TYPES = frozenset(
+    (
+        "INT64",
+        "BOOL",
+        "FLOAT64",
+        "STRING",
+        "BYTES",
+        "TIMESTAMP",
+        "DATE",
+        "TIME",
+        "DATETIME",
+        "GEOGRAPHY",
+        "NUMERIC",
+    )
+)
+
+_SQL_NONSCALAR_TYPES = frozenset(("TYPE_KIND_UNSPECIFIED", "ARRAY", "STRUCT"))
+
+
+def _make_sql_scalars_enum():
+    """Create an enum based on a gapic enum containing only SQL scalar types."""
+
+    new_enum = enum.Enum(
+        "StandardSqlDataTypes",
+        (
+            (member.name, member.value)
+            for member in gapic_enums.StandardSqlDataType.TypeKind
+            if member.name in _SQL_SCALAR_TYPES
+        ),
+    )
+
+    # make sure the docstring for the new enum is also correct
+    orig_doc = gapic_enums.StandardSqlDataType.TypeKind.__doc__
+    skip_pattern = re.compile(
+        "|".join(_SQL_NONSCALAR_TYPES)
+        + "|because a JSON object"  # the second description line of STRUCT member
+    )
+
+    new_doc = "\n".join(
+        six.moves.filterfalse(skip_pattern.search, orig_doc.splitlines())
+    )
+    new_enum.__doc__ = "An Enum of scalar SQL types.\n" + new_doc
+
+    return new_enum
+
+
+StandardSqlDataTypes = _make_sql_scalars_enum()

bigquery/tests/unit/enums/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2019, Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

bigquery/tests/unit/enums/test_standard_sql_data_types.py

@@ -0,0 +1,73 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import pytest
+
+
+@pytest.fixture
+def module_under_test():
+    from google.cloud.bigquery import enums
+
+    return enums
+
+
+@pytest.fixture
+def enum_under_test():
+    from google.cloud.bigquery.enums import StandardSqlDataTypes
+
+    return StandardSqlDataTypes
+
+
+@pytest.fixture
+def gapic_enum():
+    """The referential autogenerated enum the enum under test is based on."""
+    from google.cloud.bigquery_v2.gapic.enums import StandardSqlDataType
+
+    return StandardSqlDataType.TypeKind
+
+
+def test_all_gapic_enum_members_are_known(module_under_test, gapic_enum):
+    gapic_names = set(type_.name for type_ in gapic_enum)
+    anticipated_names = (
+        module_under_test._SQL_SCALAR_TYPES | module_under_test._SQL_NONSCALAR_TYPES
+    )
+    assert not (gapic_names - anticipated_names)  # no unhandled names
+
+
+def test_standard_sql_types_enum_members(enum_under_test, gapic_enum):
+    # check the presence of a few typical SQL types
+    for name in ("INT64", "FLOAT64", "DATE", "BOOL", "GEOGRAPHY"):
+        assert name in enum_under_test.__members__
+
+    # the enum members must match those in the original gapic enum
+    for member in enum_under_test:
+        assert member.name in gapic_enum.__members__
+        assert member.value == gapic_enum[member.name].value
+
+    # check a few members that should *not* be copied over from the gapic enum
+    for name in ("STRUCT", "ARRAY"):
+        assert name in gapic_enum.__members__
+        assert name not in enum_under_test.__members__
+
+
+def test_standard_sql_types_enum_docstring(enum_under_test, gapic_enum):
+    assert "STRUCT (int):" not in enum_under_test.__doc__
+    assert "BOOL (int):" in enum_under_test.__doc__
+    assert "TIME (int):" in enum_under_test.__doc__
+
+    # All lines in the docstring should actually come from the original docstring,
+    # except for the header.
+    assert "An Enum of scalar SQL types." in enum_under_test.__doc__
+    doc_lines = enum_under_test.__doc__.splitlines()
+    assert set(doc_lines[1:]) <= set(gapic_enum.__doc__.splitlines())