Updated init and scan to support functions.

Author: sbyrne-sf

src/datacustomcode/cli.py

@@ -145,18 +145,25 @@ def deploy(
 
 @cli.command()
 @click.argument("directory", default=".")
-def init(directory: str):
-    from datacustomcode.scan import dc_config_json_from_file
-    from datacustomcode.template import copy_template
+@click.option("--type", default="script", type=click.Choice(["script", "function"]))
+def init(directory: str, type: str):
+    from datacustomcode.scan import dc_config_json_from_file, update_config
+    from datacustomcode.template import copy_function_template, copy_script_template
 
     click.echo("Copying template to " + click.style(directory, fg="blue", bold=True))
-    copy_template(directory)
+    if type == "script":
+        copy_script_template(directory)
+    elif type == "function":
+        copy_function_template(directory)
     entrypoint_path = os.path.join(directory, "payload", "entrypoint.py")
     config_location = os.path.join(os.path.dirname(entrypoint_path), "config.json")
-    config_json = dc_config_json_from_file(entrypoint_path)
+    config_json = dc_config_json_from_file(entrypoint_path, type)
     with open(config_location, "w") as f:
         json.dump(config_json, f, indent=2)
 
+    updated_config_json = update_config(entrypoint_path)
+    with open(config_location, "w") as f:
+        json.dump(updated_config_json, f, indent=2)
     click.echo(
         "Start developing by updating the code in "
         + click.style(entrypoint_path, fg="blue", bold=True)
@@ -176,15 +183,15 @@ def init(directory: str):
     "--no-requirements", is_flag=True, help="Skip generating requirements.txt file"
 )
 def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
-    from datacustomcode.scan import dc_config_json_from_file, write_requirements_file
+    from datacustomcode.scan import update_config, write_requirements_file
 
     config_location = config or os.path.join(os.path.dirname(filename), "config.json")
     click.echo(
         "Dumping scan results to config file: "
         + click.style(config_location, fg="blue", bold=True)
     )
     click.echo("Scanning " + click.style(filename, fg="blue", bold=True) + "...")
-    config_json = dc_config_json_from_file(filename)
+    config_json = update_config(filename)
 
     click.secho(json.dumps(config_json, indent=2), fg="yellow")
     if not dry_run:

src/datacustomcode/scan.py

@@ -37,14 +37,20 @@
 
 DATA_TRANSFORM_CONFIG_TEMPLATE = {
     "sdkVersion": get_version(),
+    "type": "script",
     "entryPoint": "",
-    "dataspace": "",
+    "dataspace": "default",
     "permissions": {
         "read": {},
         "write": {},
     },
 }
 
+FUNCTION_CONFIG_TEMPLATE = {
+    "sdkVersion": get_version(),
+    "type": "function",
+    "entryPoint": "",
+}
 STANDARD_LIBS = set(sys.stdlib_module_names)
 
 
@@ -230,57 +236,66 @@ def scan_file(file_path: str) -> DataAccessLayerCalls:
         return visitor.found()
 
 
-def dc_config_json_from_file(file_path: str) -> dict[str, Any]:
+def dc_config_json_from_file(file_path: str, type: str) -> dict[str, Any]:
     """Create a Data Cloud Custom Code config JSON from a script."""
-    output = scan_file(file_path)
-    config = DATA_TRANSFORM_CONFIG_TEMPLATE.copy()
+    config: dict[str, Any]
+    if type == "script":
+        config = DATA_TRANSFORM_CONFIG_TEMPLATE.copy()
+    elif type == "function":
+        config = FUNCTION_CONFIG_TEMPLATE.copy()
     config["entryPoint"] = file_path.rpartition("/")[-1]
+    return config
 
+
+def update_config(file_path: str) -> dict[str, Any]:
     file_dir = os.path.dirname(file_path)
     config_json_path = os.path.join(file_dir, "config.json")
-
+    existing_config: dict[str, Any]
     if os.path.exists(config_json_path) and os.path.isfile(config_json_path):
         try:
             with open(config_json_path, "r") as f:
                 existing_config = json.load(f)
-
-            if "dataspace" in existing_config:
-                dataspace_value = existing_config["dataspace"]
-                if not dataspace_value or (
-                    isinstance(dataspace_value, str) and dataspace_value.strip() == ""
-                ):
-                    logger.warning(
-                        f"dataspace in {config_json_path} is empty or None. "
-                        f"Updating config file to use dataspace 'default'. "
-                    )
-                    config["dataspace"] = "default"
-                else:
-                    config["dataspace"] = dataspace_value
-            else:
-                raise ValueError(
-                    f"dataspace must be defined in {config_json_path}. "
-                    f"Please add a 'dataspace' field to the config.json file. "
-                )
         except json.JSONDecodeError as e:
             raise ValueError(
                 f"Failed to parse JSON from {config_json_path}: {e}"
             ) from e
         except OSError as e:
             raise OSError(f"Failed to read config file {config_json_path}: {e}") from e
     else:
-        config["dataspace"] = "default"
-
-    read: dict[str, list[str]] = {}
-    if output.read_dlo:
-        read["dlo"] = list(output.read_dlo)
-    else:
-        read["dmo"] = list(output.read_dmo)
-    write: dict[str, list[str]] = {}
-    if output.write_to_dlo:
-        write["dlo"] = list(output.write_to_dlo)
+        raise ValueError(f"config.json not found at {config_json_path}")
+    if existing_config["type"] == "script":
+        existing_config["dataspace"] = get_dataspace(existing_config)
+        output = scan_file(file_path)
+        read: dict[str, list[str]] = {}
+        if output.read_dlo:
+            read["dlo"] = list(output.read_dlo)
+        else:
+            read["dmo"] = list(output.read_dmo)
+        write: dict[str, list[str]] = {}
+        if output.write_to_dlo:
+            write["dlo"] = list(output.write_to_dlo)
+        else:
+            write["dmo"] = list(output.write_to_dmo)
+
+        existing_config["permissions"] = {"read": read, "write": write}
+    return existing_config
+
+
+def get_dataspace(existing_config: dict[str, str]) -> str:
+    if "dataspace" in existing_config:
+        dataspace_value = existing_config["dataspace"]
+        if not dataspace_value or (
+            isinstance(dataspace_value, str) and dataspace_value.strip() == ""
+        ):
+            logger.warning(
+                "dataspace is empty or None. "
+                "Updating config file to use dataspace 'default'. "
+            )
+            return "default"
+        else:
+            return dataspace_value
     else:
-        write["dmo"] = list(output.write_to_dmo)
-
-    config["permissions"] = {"read": read, "write": write}
-
-    return config
+        raise ValueError(
+            "dataspace must be defined. "
+            "Please add a 'dataspace' field to the config.json file. "
+        )

src/datacustomcode/template.py

@@ -17,15 +17,31 @@
 
 from loguru import logger
 
-template_dir = os.path.join(os.path.dirname(__file__), "templates")
+script_template_dir = os.path.join(os.path.dirname(__file__), "templates", "script")
+function_template_dir = os.path.join(os.path.dirname(__file__), "templates", "function")
 
 
-def copy_template(target_dir: str) -> None:
+def copy_script_template(target_dir: str) -> None:
     """Copy the template to the target directory."""
     os.makedirs(target_dir, exist_ok=True)
 
-    for item in os.listdir(template_dir):
-        source = os.path.join(template_dir, item)
+    for item in os.listdir(script_template_dir):
+        source = os.path.join(script_template_dir, item)
+        destination = os.path.join(target_dir, item)
+
+        if os.path.isdir(source):
+            logger.debug(f"Copying directory {source} to {destination}...")
+            shutil.copytree(source, destination, dirs_exist_ok=True)
+        else:
+            logger.debug(f"Copying file {source} to {destination}...")
+            shutil.copy2(source, destination)
+
+
+def copy_function_template(target_dir: str) -> None:
+    os.makedirs(target_dir, exist_ok=True)
+
+    for item in os.listdir(function_template_dir):
+        source = os.path.join(function_template_dir, item)
         destination = os.path.join(target_dir, item)
 
         if os.path.isdir(source):