create template files

Author: tswast

packages/bigframes/bigframes/bigquery/_googlesql.py

@@ -58,7 +58,7 @@ def apply_googlesql_scalar_op(
             if isinstance(arg, bigframes.core.col.Expression):
                 block, col_id = block.project_expr(bigframes.core.col._as_bf_expr(arg))
                 processed_args.append(series.Series(block.select_column(col_id)))
-            elif arg is sentinels.DEFAULT:
+            elif arg is sentinels.Sentinel.ARGUMENT_DEFAULT:
                 processed_args.append(bigframes.core.col.Expression(ex.OmittedArg()))
             else:
                 processed_args.append(arg)
@@ -73,7 +73,7 @@ def apply_googlesql_scalar_op(
     for arg in args:
         if isinstance(arg, bigframes.core.col.Expression):
             expr_args.append(bigframes.core.col._as_bf_expr(arg))
-        elif arg is sentinels.DEFAULT:
+        elif arg is sentinels.Sentinel.ARGUMENT_DEFAULT:
             expr_args.append(ex.OmittedArg())
         else:
             expr_args.append(ex.const(arg))

packages/bigframes/bigframes/bigquery/_operations/aead.py

@@ -11,7 +11,6 @@
 # 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.
-
 #
 # DO NOT MODIFY THIS FILE DIRECTLY.
 # This file was generated from: scripts/data/sql-functions/aead.yaml
@@ -20,7 +19,7 @@
 from __future__ import annotations
 
 import datetime
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, Literal, Optional, TypeVar, Union
 
 import bigframes.bigquery._googlesql
 import bigframes.core.col
@@ -51,9 +50,9 @@
 
 
 def decrypt_bytes(
-    keyset: Union[T, bigframes.core.col.Expression, Union[bytes, dict]],
-    ciphertext: Union[T, bigframes.core.col.Expression, bytes],
-    additional_data: Union[T, bigframes.core.col.Expression, bytes],
+    keyset: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict]],
+    ciphertext: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes]],
+    additional_data: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes]],
 ) -> T:
     """Uses the matching key from keyset to decrypt ciphertext and verifies the integrity of the data using additional_data. Returns an error if decryption or verification fails."""
     return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
@@ -65,9 +64,9 @@ def decrypt_bytes(
 
 
 def decrypt_string(
-    keyset: Union[T, bigframes.core.col.Expression, Union[bytes, dict]],
-    ciphertext: Union[T, bigframes.core.col.Expression, bytes],
-    additional_data: Union[T, bigframes.core.col.Expression, str],
+    keyset: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict]],
+    ciphertext: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes]],
+    additional_data: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], str]],
 ) -> T:
     """Like AEAD.DECRYPT_BYTES, but where additional_data is of type STRING."""
     return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
@@ -79,9 +78,9 @@ def decrypt_string(
 
 
 def encrypt(
-    keyset: Union[T, bigframes.core.col.Expression, Union[bytes, dict]],
-    plaintext: Union[T, bigframes.core.col.Expression, Union[bytes, str]],
-    additional_data: Union[T, bigframes.core.col.Expression, Union[bytes, str]],
+    keyset: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict]],
+    plaintext: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str]],
+    additional_data: Union[T, bigframes.core.col.Expression, Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str]],
 ) -> T:
     """Encrypts plaintext using the primary cryptographic key in keyset. The algorithm of the primary key must be AEAD_AES_GCM_256. Binds the ciphertext to the context defined by additional_data. Returns NULL if any input is NULL."""
     return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(

packages/bigframes/bigframes/core/sentinels.py

@@ -16,18 +16,16 @@
 
 from __future__ import annotations
 
-from enum import Enum
+import enum
 
 
-class Default(Enum):
-    """Default values used throughout BigFrames.
+class Sentinel(enum.Enum):
+    """Default values used throughout BigFrames."""
+
+    """Default value for an optional argument.
 
     When a parameter is set to this, that parameter is explicitly omitted
     from the SQL text. This allows for NULL (None in Python) to be explicitly
     passed in to optional parameters.
     """
-
-    token = 0
-
-
-DEFAULT = Default.token
+    ARGUMENT_DEFAULT = enum.auto()

packages/bigframes/scripts/generate_bigframes_bigquery.py

@@ -35,6 +35,8 @@
 OUTPUT_DIR = pathlib.Path("bigframes/bigquery/_operations")
 # Directory where the generated test files will be placed
 TEST_OUTPUT_DIR = pathlib.Path("tests/unit/bigquery/_operations")
+# Directory containing the Jinja2 templates
+TEMPLATE_DIR = pathlib.Path("scripts/templates")
 
 RUFF_ARGS = [
     "ruff",
@@ -46,97 +48,6 @@
     "--line-length=88",
 ]
 
-LICENSE_HEADER = """# Copyright 2026 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.
-"""
-
-TEMPLATE = """{{ license_header }}
-#
-# DO NOT MODIFY THIS FILE DIRECTLY.
-# This file was generated from: {{ yaml_path }}
-# by the script: {{ script_path }}
-
-from __future__ import annotations
-
-import datetime
-from typing import Any, Optional, TypeVar, Union
-
-from bigframes import dtypes
-import bigframes.bigquery._googlesql
-import bigframes.core.col
-import bigframes.core.expression as ex
-import bigframes.core.sentinels as sentinels
-from bigframes.operations import googlesql
-import bigframes.operations as ops
-import bigframes.series as series
-
-T = TypeVar("T", series.Series, bigframes.core.col.Expression)
-
-{% for op in ops %}
-{{ op.internal_name }} = googlesql.GoogleSqlScalarOp(
-    "{{ op.sql_name }}",
-    args=({{ op.arg_specs }}),
-    signature={{ op.signature }},
-)
-{% endfor %}
-{% for func in functions %}
-
-
-def {{ func.name }}(
-{% for arg in func.args %}
-    {{ arg.name }}: Union[T, bigframes.core.col.Expression, {{ arg.type_hint }}]{% if arg.default %} = {{ arg.default }}{% endif %},
-{% endfor %}
-) -> T:
-    \"\"\"{{ func.description }}\"\"\"
-    return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
-        {{ func.op_name }},
-{% for arg in func.args %}
-        {{ arg.name }},
-{% endfor %}
-    )  # type: ignore
-{% endfor %}
-"""
-
-TEST_TEMPLATE = r"""{{ license_header }}
-#
-# DO NOT MODIFY THIS FILE DIRECTLY.
-# This file was generated from: {{ yaml_path }}
-# by the script: {{ script_path }}
-
-from typing import cast
-
-import pytest
-
-import bigframes.pandas as bpd
-import {{ import_path }} as {{ short_name }}
-
-pytest.importorskip("pytest_snapshot")
-
-
-{% for func in functions %}
-def test_{{ func.name }}(scalar_types_df: bpd.DataFrame, snapshot):
-    result = {{ short_name }}.{{ func.name }}(
-{% for arg in func.test_args %}
-        cast(bpd.Series, scalar_types_df["{{ arg.col_name }}"]),
-{% endfor %}
-    ).to_frame()
-    snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql")
-
-
-{% endfor %}
-"""
-
 DTYPE_MAP = {
     "binary": "dtypes.BYTES_DTYPE",
     "string": "dtypes.STRING_DTYPE",
@@ -191,9 +102,13 @@ def to_snake_case(name):
 
 
 def main():
-    env = jinja2.Environment(trim_blocks=True, lstrip_blocks=True)
-    template = env.from_string(TEMPLATE)
-    test_template = env.from_string(TEST_TEMPLATE)
+    env = jinja2.Environment(
+        loader=jinja2.FileSystemLoader(TEMPLATE_DIR),
+        trim_blocks=True,
+        lstrip_blocks=True,
+    )
+    template = env.get_template("operation.py.j2")
+    test_template = env.get_template("test_operation.py.j2")
 
     for yaml_file in DATA_DIR.glob("**/*.yaml"):
         print(f"Processing {yaml_file}...")
@@ -266,13 +181,13 @@ def main():
                 func_args = []
                 for name in arg_order:
                     arg_info = args_by_name[name]
-                    types = [PY_TYPE_MAP.get(t, "Any") for t in arg_info["types"]]
+                    types = [PY_TYPE_MAP.get(t, "Any") for t in arg_info["types"]] + ["Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]"]
                     type_hint = (
                         "Union[" + ", ".join(sorted(set(types))) + "]"
                         if len(types) > 1
                         else types[0]
                     )
-                    default = "sentinels.DEFAULT" if arg_info["optional"] else ""
+                    default = "sentinels.Sentinel.ARGUMENT_DEFAULT" if arg_info["optional"] else ""
                     func_args.append(
                         {
                             "name": name,
@@ -308,7 +223,6 @@ def main():
         # Render and write
         output_file.parent.mkdir(parents=True, exist_ok=True)
         content = template.render(
-            license_header=LICENSE_HEADER,
             yaml_path=str(yaml_file),
             script_path="scripts/generate_bigframes_bigquery.py",
             ops=ops_list,
@@ -334,7 +248,6 @@ def main():
 
         test_output_file.parent.mkdir(parents=True, exist_ok=True)
         test_content = test_template.render(
-            license_header=LICENSE_HEADER,
             yaml_path=str(yaml_file),
             script_path="scripts/generate_bigframes_bigquery.py",
             import_path=import_path,

packages/bigframes/scripts/templates/license.py.j2

@@ -0,0 +1,13 @@
+# Copyright 2026 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.

packages/bigframes/scripts/templates/operation.py.j2

@@ -0,0 +1,46 @@
+{% include 'license.py.j2' %}
+
+#
+# DO NOT MODIFY THIS FILE DIRECTLY.
+# This file was generated from: {{ yaml_path }}
+# by the script: {{ script_path }}
+
+from __future__ import annotations
+
+import datetime
+from typing import Any, Literal, Optional, TypeVar, Union
+
+from bigframes import dtypes
+import bigframes.bigquery._googlesql
+import bigframes.core.col
+import bigframes.core.expression as ex
+import bigframes.core.sentinels as sentinels
+from bigframes.operations import googlesql
+import bigframes.operations as ops
+import bigframes.series as series
+
+T = TypeVar("T", series.Series, bigframes.core.col.Expression)
+
+{% for op in ops %}
+{{ op.internal_name }} = googlesql.GoogleSqlScalarOp(
+    "{{ op.sql_name }}",
+    args=({{ op.arg_specs }}),
+    signature={{ op.signature }},
+)
+{% endfor %}
+{% for func in functions %}
+
+
+def {{ func.name }}(
+{% for arg in func.args %}
+    {{ arg.name }}: Union[T, bigframes.core.col.Expression, {{ arg.type_hint }}]{% if arg.default %} = {{ arg.default }}{% endif %},
+{% endfor %}
+) -> T:
+    """{{ func.description }}"""
+    return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
+        {{ func.op_name }},
+{% for arg in func.args %}
+        {{ arg.name }},
+{% endfor %}
+    )  # type: ignore
+{% endfor %}

packages/bigframes/scripts/templates/test_operation.py.j2

@@ -0,0 +1,28 @@
+{% include 'license.py.j2' %}
+
+#
+# DO NOT MODIFY THIS FILE DIRECTLY.
+# This file was generated from: {{ yaml_path }}
+# by the script: {{ script_path }}
+
+from typing import cast
+
+import pytest
+
+import bigframes.pandas as bpd
+import {{ import_path }} as {{ short_name }}
+
+pytest.importorskip("pytest_snapshot")
+
+
+{% for func in functions %}
+def test_{{ func.name }}(scalar_types_df: bpd.DataFrame, snapshot):
+    result = {{ short_name }}.{{ func.name }}(
+{% for arg in func.test_args %}
+        cast(bpd.Series, scalar_types_df["{{ arg.col_name }}"]),
+{% endfor %}
+    ).to_frame()
+    snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql")
+
+
+{% endfor %}