checkpoint: impl operator definition gen

Author: sycai

packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py

@@ -22,18 +22,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pprint
 
 import constants
-import yaml_parser
 import file_generator
+import yaml_parser
 
 
 def main():
     modules = []
 
     for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")):
         modules.append(yaml_parser.parse_yaml(yaml_file))
+        break
 
     file_generator.generate(modules)
 

packages/bigframes/scripts/generate_bigframes_bigquery/constants.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 
 import pathlib
+
 import jinja2
 
 SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute()
 PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent
 CODE_ROOT = PACKAGE_ROOT / "bigframes"
-SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT)
+SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent
 
 # Directory containing the YAML files
 DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions"
@@ -119,6 +120,59 @@
     "decimal<76,38>": "dtypes.BIGNUMERIC_DTYPE",
 }
 
+YAML_TYPE_TO_COL = {
+    "binary": "bytes_col",
+    "string": "string_col",
+    "int64": "int64_col",
+    "i64": "int64_col",
+    "float64": "float64_col",
+    "fp64": "float64_col",
+    "bool": "bool_col",
+    "boolean": "bool_col",
+    "geography": "geography_col",
+    "date": "date_col",
+    "time": "time_col",
+    "datetime": "datetime_col",
+    "timestamp": "timestamp_col",
+    "decimal<38,9>": "numeric_col",
+    "decimal<76,38>": "bignumeric_col",
+}
+
+PY_TYPE_MAP = {
+    "binary": "bytes",
+    "string": "str",
+    "int64": "int",
+    "i64": "int",
+    "float64": "float",
+    "fp64": "float",
+    "bool": "bool",
+    "boolean": "bool",
+    "geography": "Any",
+    "json": "Any",
+    "date": "datetime.date",
+    "time": "datetime.time",
+    "datetime": "datetime.datetime",
+    "timestamp": "datetime.datetime",
+    "struct": "dict",
+    "decimal<38,9>": "decimal.Decimal",
+    "decimal<76,38>": "decimal.Decimal",
+    "interval_day": "datetime.timedelta",
+}
+
+RUFF_COMMON_ARGS = [
+    "--target-version=py310",
+    "--line-length=88",
+]
+RUFF_CHECK_ARGS = [
+    "check",
+    "--select",
+    "I,F",
+    "--fix",
+] + RUFF_COMMON_ARGS
+RUFF_FORMAT_ARGS = [
+    "format",
+] + RUFF_COMMON_ARGS
+
 
 TEMPLATES: dict[str, jinja2.Template]  # Lazily-loaded global variable
 _templates = None

packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py

@@ -15,9 +15,9 @@
 
 import dataclasses
 import pathlib
+from typing import Any
 
 import constants
-from typing import Any
 
 
 @dataclasses.dataclass
@@ -28,6 +28,14 @@ class BQFuncArg:
     keyword_only: bool
 
 
+@dataclasses.dataclass
+class BigFramesFuncArg:
+    name: str
+    types: set[str]
+    optional: bool
+    keyword_only: bool
+
+
 @dataclasses.dataclass
 class BQFuncImpl:
     args: list[BQFuncArg]
@@ -56,26 +64,19 @@ def to_dict(self) -> dict[str, Any]:
 @dataclasses.dataclass
 class BQFunc:
     name: str
+    op_base_name: str
     description: str
     impls: list[BQFuncImpl]
 
 
 @dataclasses.dataclass
 class BQModule:
-    module_path: pathlib.Path
+    yaml_file: pathlib.Path
     functions: list[BQFunc]
 
     @property
-    def name(self) -> str:
-        return self.module_path.name
-
-    @property
-    def is_global_namespace(self) -> bool:
-        return "global_namespace" in self.module_path.parts
-
-    @property
-    def output_file(self):
-        return constants.OUTPUT_DIR.joinpath(self.module_path).with_suffix(".py")
+    def module_path(self):
+        return self.yaml_file.relative_to(constants.DATA_DIR).with_suffix("")
 
 
 @dataclasses.dataclass
@@ -84,12 +85,13 @@ class BigFramesOp:
     sql_name: str
     arg_specs: str
     signature: str
-    signature_definition: str
+    signature_definition: str | None
 
 
 @dataclasses.dataclass
 class BigFramesFunc:
     name: str
+    op_name: str
     description: str
     args: list[str]
     series_accessor_arg: list[str]

packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py

@@ -12,33 +12,76 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pathlib
+import sys
+import subprocess
+
 import constants
-import jinja2
 import data_models
-import pprint
-
 import template_renderer
 
 
-def _generate_sql_operator_def():
-    pass
+def _ensure_init_py(directory: pathlib.Path, limit_dir: pathlib.Path):
+    """Ensures __init__.py exists in the directory and its parents up to limit_dir."""
+    curr = directory
+    while curr != limit_dir and curr != curr.parent:
+        init_file = curr / "__init__.py"
+        if not init_file.exists():
+            print(f"  Creating {init_file}")
+            content = constants.TEMPLATES["license"].render()
+            with open(init_file, "w") as f:
+                f.write(content)
+        curr = curr.parent
+
+
+def _write_file(content: str, output_file: pathlib.Path, limit_dir: pathlib.Path):
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    _ensure_init_py(output_file.parent, limit_dir)
+
+    with open(output_file, "w") as f:
+        f.write(content)
+    print(f"  Generated {output_file}")
+
+
+def _run_ruff():
+    targets = [
+        constants.OUTPUT_DIR,
+        constants.TEST_OUTPUT_DIR,
+    ]
+
+    subprocess.run(
+        [sys.executable, "-m", "ruff"] + constants.RUFF_CHECK_ARGS + targets,
+        check=True,
+    )
+
+    subprocess.run(
+        [sys.executable, "-m", "ruff"] + constants.RUFF_FORMAT_ARGS + targets,
+        check=True,
+    )
+
+
+def _generate_op_defs(bq_module: data_models.BQModule):
+    content = template_renderer.render_operation(bq_module)
+
+    output_file = constants.OUTPUT_DIR.joinpath(bq_module.module_path).with_suffix(
+        ".py"
+    )
+
+    _write_file(content, output_file, constants.OUTPUT_DIR.parent)
 
 
 def _generate_accesor():
     pass
 
 
-def _generate_tests():
+def _generate_tests(bq_module: data_models.BQModule):
     pass
 
 
 def generate(bq_modules: list[data_models.BQModule]):
 
     for bq_module in bq_modules:
-        for bq_func in bq_module.functions:
-            pprint.pp(template_renderer.render_signature_def(bq_func, bq_module))
-
-    # Write to file
+        _generate_op_defs(bq_module)
 
     # Ruff format
-    pass
+    _run_ruff()