gen type asserts in python

Signed-off-by: Isabella do Amaral <[email protected]>

scripts/gen_type_asserts.py

@@ -0,0 +1,70 @@
+import collections.abc as c
+from itertools import islice
+from pathlib import Path
+from sys import argv
+from textwrap import dedent
+
+
+def get_funcs(models: c.Iterable[Path]) -> c.Iterator[str]:
+    for path in models:
+        with path.open() as f:
+            buf = []
+            in_func = False
+            # skip first 14 lines of boilerplate
+            for raw in islice(f.readlines(), 14, None):
+                line = raw.rstrip()
+                buf.append(line)
+                if line.startswith("func"):
+                    in_func = True
+                elif line.startswith("}") and in_func:
+                    in_func = False
+                    yield "\n".join(buf)
+                    buf.clear()
+        path.unlink()
+
+
+def get_name(func: str) -> str:
+    # each func is declared as:
+    # // comment
+    # func funcName(args) error ...
+    return func.split("\n")[1].split(" ")[1].split("(")[0]
+
+
+if __name__ == "__main__":
+    if len(argv) > 1:
+        root = Path(argv[1])
+        assert root.is_dir(), f"{root} is not a directory"
+    else:
+        # model registry repo root
+        root = Path(__file__).parent.parent.resolve()
+    src = root / "internal/server/openapi"
+    print(
+        dedent("""
+/*
+ * Model Registry REST API
+ *
+ * REST API for Model Registry to create and manage ML model metadata
+ *
+ * API version: 1.0.0
+ * Generated by: OpenAPI Generator (https://openapi-generator.tech)
+ *
+ */
+
+// File generated by scripts/gen_type_assert.sh - DO NOT EDIT
+
+package openapi
+
+import (
+    model "github.com/kubeflow/model-registry/pkg/openapi"
+)
+
+    """)
+    )
+    print(
+        "\n".join(
+            sorted(
+                get_funcs(path for path in src.rglob("model_*.go")),
+                key=lambda f: get_name(f),
+            )
+        )
+    )

scripts/gen_type_asserts.sh

@@ -7,78 +7,7 @@ PROJECT_ROOT=$(realpath "$(dirname "$0")"/..)
 ASSERT_FILE_PATH="${PROJECT_ROOT}/internal/server/openapi/type_asserts.go"
 PATCH="${PROJECT_ROOT}/patches/type_asserts.patch"
 
-# Remove the existing file identified by env.ASSERT_FILE_PATH
-if [ -f "$ASSERT_FILE_PATH" ]; then
-    rm "$ASSERT_FILE_PATH"
-fi
-
-# Create an empty file
-touch "$ASSERT_FILE_PATH"
-
-INITIAL_CONTENT=$(
-    cat <<EOF
-/*
- * Model Registry REST API
- *
- * REST API for Model Registry to create and manage ML model metadata
- *
- * API version: 1.0.0
- * Generated by: OpenAPI Generator (https://openapi-generator.tech)
- *
- */
-
-// File generated by scripts/gen_type_assert.sh - DO NOT EDIT
-
-package openapi
-
-import (
-	model "github.com/kubeflow/model-registry/pkg/openapi"
-)
-
-
-EOF
-)
-
-# Create the file and initialize it with the specified content
-echo -e "$INITIAL_CONTENT" >"$ASSERT_FILE_PATH"
-
-declare -A assert_functions
-pattern="\/\/ (\w+) checks(.|\n)+?\n\}\n"
-shell_pattern="\/\/ (\w+) checks[^\/?]+?"
-
-# Iterate over files starting with "model_" in the internal/server/openapi/ folder
-for file in "$PROJECT_ROOT"/internal/server/openapi/model_*; do
-    # Check if the file is a regular file
-    if [ -f "$file" ]; then
-
-        # grab all functions in the file
-        functions=$(grep -Pzo "$pattern" "${file}" | tr -d '\0')
-        while [[ $functions =~ $shell_pattern ]]; do
-
-          name="${BASH_REMATCH[1]}"
-          body="${BASH_REMATCH[0]}"
-
-          # add function to associative array
-          assert_functions["$name"]="$body"
-
-          # Remove the matched function to process next function
-          functions=${functions//"$body"/}
-        done
-
-        # Remove the merged file
-        rm "$file"
-    fi
-done
-
-# get sorted function names in another array
-sorted_names=$(for name in "${!assert_functions[@]}"; do
-  echo "$name"
-done | sort)
-
-# iterate over the sorted function names array and print bodies
-for name in $sorted_names; do
-  echo "${assert_functions[$name]}" >>"$ASSERT_FILE_PATH"
-done
+python "${PROJECT_ROOT}/scripts/gen_type_asserts.py" >"$ASSERT_FILE_PATH"
 
 gofmt -w "$ASSERT_FILE_PATH"