Zip and lint cleanup

Author: markdlv-sf

src/datacustomcode/cli.py

@@ -68,34 +68,14 @@ def configure(
         login_url=login_url,
     ).update_ini(profile=profile)
 
+
 @cli.command()
-@click.option("--profile", default="default")
-@click.option("--path", default="payload")
-@click.option("--name", default="test_pkg")
-@click.option("--version", default="0.0.1")
-@click.option("--description", default="Custom Data Transform Code")
-def zip(profile: str, path: str, name: str, version: str, description: str):
-    from datacustomcode.credentials import Credentials
-    from datacustomcode.deploy import TransformationJobMetadata, zip, zip_and_upload_directory
+@click.argument("path", default="payload")
+def zip(path: str):
+    from datacustomcode.deploy import zip
 
     logger.debug("Zipping project")
-
-    metadata = TransformationJobMetadata(
-        name=name,
-        version=version,
-        description=description,
-    )
-    try:
-        credentials = Credentials.from_ini(profile=profile)
-    except KeyError:
-        click.secho(
-            f"Error: Profile {profile} not found in credentials.ini. "
-            "Run `datacustomcode configure` to create a credentialsprofile.",
-            fg="red",
-        )
-        raise click.Abort() from None
-    zip(path, metadata, credentials, name)
-
+    zip(path)
 
 
 @cli.command()
@@ -156,11 +136,12 @@ def init(directory: str):
 @click.argument("filename")
 @click.option("--config")
 @click.option("--dry-run", is_flag=True)
-@click.option("--no-requirements", is_flag=True, help="Skip generating requirements.txt file")
+@click.option(
+    "--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
 
-
     config_location = config or os.path.join(os.path.dirname(filename), "config.json")
     click.echo(
         "Dumping scan results to config file: "

src/datacustomcode/deploy.py

@@ -177,35 +177,6 @@ def prepare_dependency_archive(directory: str) -> None:
         logger.debug(f"Dependencies downloaded and archived to {archive_file}")
 
 
-def zip_and_upload_directory(directory: str, name: str) -> None:
-    #   file_upload_url = unescape(file_upload_url)
-
-    logger.debug(f"Zipping directory... {directory}")
-
-    # Create a zip file excluding .DS_Store files
-    import zipfile
-
-    zip_filename = f"{name}.zip"
-    with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, dirs, files in os.walk(directory):
-            # Skip .DS_Store files when adding to zip
-            for file in files:
-                if file != '.DS_Store':
-                    file_path = os.path.join(root, file)
-                    # Preserve relative path structure in the zip file
-                    arcname = os.path.relpath(file_path, directory)
-                    zipf.write(file_path, arcname)
-
-    logger.debug(f"Created zip file: {zip_filename} (excluding .DS_Store files)")
-
-    # logger.debug(f"Uploading deployment to {file_upload_url}")
-    # with open(ZIP_FILE_NAME, "rb") as zip_file:
-    #     response = requests.put(
-    #         file_upload_url, data=zip_file, headers={"Content-Type": "application/zip"}
-    #     )
-    #     response.raise_for_status()
-
-
 class DeploymentsResponse(BaseModel):
     deploymentStatus: str
 
@@ -342,59 +313,70 @@ def create_data_transform(
     response = _make_api_call(url, "POST", token=access_token.access_token, json=body)
     return response
 
+
 def has_nonempty_requirements_file(directory: str) -> bool:
     """
-    Check if requirements.txt exists in the given directory and has at least one non-comment line.
+    Check if requirements.txt exists in the given directory and has at least
+    one non-comment line.
     Args:
         directory (str): The directory to check for requirements.txt.
     Returns:
-        bool: True if requirements.txt exists and has a non-comment line, False otherwise.
+        bool: True if requirements.txt exists and has a non-comment line,
+        False otherwise.
     """
     # Look for requirements.txt in the parent directory of the given directory
     requirements_path = os.path.join(os.path.dirname(directory), "requirements.txt")
-    print(requirements_path)
 
     try:
         if os.path.isfile(requirements_path):
-            #print the contents of the file
-            with open(requirements_path, "r", encoding="utf-8") as f:
-                print(f.read())
             with open(requirements_path, "r", encoding="utf-8") as f:
                 for line in f:
-                    # Consider non-empty if any line is not a comment (ignoring leading whitespace)
-                    if line.strip() and not line.lstrip().startswith('#'):
+                    # Consider non-empty if any line is not a comment (ignoring
+                    # leading whitespace)
+                    if line.strip() and not line.lstrip().startswith("#"):
                         return True
     except Exception as e:
         logger.error(f"Error reading requirements.txt: {e}")
     return False
 
 
+def upload_zip(file_upload_url: str) -> None:
+    file_upload_url = unescape(file_upload_url)
+    with open(ZIP_FILE_NAME, "rb") as zip_file:
+        response = requests.put(
+            file_upload_url, data=zip_file, headers={"Content-Type": "application/zip"}
+        )
+        response.raise_for_status()
+
+
 def zip(
     directory: str,
-    metadata: TransformationJobMetadata,
-    credentials: Credentials,
-    name: str,
-    callback=None,
-) -> AccessTokenResponse:
-    """Deploy a data transform in the DataCloud."""
-    access_token = _retrieve_access_token(credentials)
+):
+    # Create a zip file excluding .DS_Store files
+    import zipfile
 
     # prepare payload only if requirements.txt is non-empty
     if has_nonempty_requirements_file(directory):
         prepare_dependency_archive(directory)
     else:
-        logger.info(f"Skipping dependency archive: requirements.txt is missing or empty in {directory}")
- #   create_data_transform_config(directory)
+        logger.info(
+            f"Skipping dependency archive: requirements.txt is missing or empty "
+            f"in {directory}"
+        )
 
-    # create deployment and upload payload
-#    deployment = create_deployment(access_token, metadata)
-    zip_and_upload_directory(directory, name)
-    #, deployment.fileUploadUrl)
-#    wait_for_deployment(access_token, metadata, callback)
+    logger.debug(f"Zipping directory... {directory}")
 
-    # create data transform
-#   create_data_transform(directory, access_token, metadata)
-    return access_token
+    with zipfile.ZipFile(ZIP_FILE_NAME, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, dirs, files in os.walk(directory):
+            # Skip .DS_Store files when adding to zip
+            for file in files:
+                if file != ".DS_Store":
+                    file_path = os.path.join(root, file)
+                    # Preserve relative path structure in the zip file
+                    arcname = os.path.relpath(file_path, directory)
+                    zipf.write(file_path, arcname)
+
+    logger.debug(f"Created zip file: {ZIP_FILE_NAME}")
 
 
 def deploy_full(
@@ -412,7 +394,8 @@ def deploy_full(
 
     # create deployment and upload payload
     deployment = create_deployment(access_token, metadata)
-    zip_and_upload_directory(directory, deployment.fileUploadUrl)
+    zip(directory)
+    upload_zip(deployment.fileUploadUrl)
     wait_for_deployment(access_token, metadata, callback)
 
     # create data transform

src/datacustomcode/scan.py

@@ -13,13 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import annotations
-import os
-from typing import Union, Dict, List, FrozenSet, Set
 
 import ast
+import os
 from typing import (
     Any,
+    ClassVar,
     Dict,
+    Set,
     Union,
 )
 
@@ -137,17 +138,55 @@ class ImportVisitor(ast.NodeVisitor):
     """AST Visitor that extracts external package imports from Python code."""
 
     # Standard library modules that should be excluded from requirements
-    STANDARD_LIBS = {
-        "abc", "argparse", "ast", "asyncio", "base64", "collections", "configparser",
-        "contextlib", "copy", "csv", "datetime", "enum", "functools", "glob", "hashlib",
-        "http", "importlib", "inspect", "io", "itertools", "json", "logging", "math",
-        "os", "pathlib", "pickle", "random", "re", "shutil", "site", "socket", "sqlite3",
-        "string", "subprocess", "sys", "tempfile", "threading", "time", "traceback",
-        "typing", "uuid", "warnings", "xml", "zipfile"
+    STANDARD_LIBS: ClassVar[set[str]] = {
+        "abc",
+        "argparse",
+        "ast",
+        "asyncio",
+        "base64",
+        "collections",
+        "configparser",
+        "contextlib",
+        "copy",
+        "csv",
+        "datetime",
+        "enum",
+        "functools",
+        "glob",
+        "hashlib",
+        "http",
+        "importlib",
+        "inspect",
+        "io",
+        "itertools",
+        "json",
+        "logging",
+        "math",
+        "os",
+        "pathlib",
+        "pickle",
+        "random",
+        "re",
+        "shutil",
+        "site",
+        "socket",
+        "sqlite3",
+        "string",
+        "subprocess",
+        "sys",
+        "tempfile",
+        "threading",
+        "time",
+        "traceback",
+        "typing",
+        "uuid",
+        "warnings",
+        "xml",
+        "zipfile",
     }
 
     # Additional packages to exclude from requirements.txt
-    EXCLUDED_PACKAGES = {
+    EXCLUDED_PACKAGES: ClassVar[set[str]] = {
         "datacustomcode",  # Internal package
         "pyspark",  # Provided by the runtime environment
     }
@@ -159,21 +198,25 @@ def visit_Import(self, node: ast.Import) -> None:
         """Visit an import statement (e.g., import os, sys)."""
         for name in node.names:
             # Get the top-level package name
-            package = name.name.split('.')[0]
-            if (package not in self.STANDARD_LIBS and
-                    package not in self.EXCLUDED_PACKAGES and
-                    not package.startswith('_')):
+            package = name.name.split(".")[0]
+            if (
+                package not in self.STANDARD_LIBS
+                and package not in self.EXCLUDED_PACKAGES
+                and not package.startswith("_")
+            ):
                 self.imports.add(package)
         self.generic_visit(node)
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
         """Visit a from-import statement (e.g., from os import path)."""
         if node.module is not None:
             # Get the top-level package
-            package = node.module.split('.')[0]
-            if (package not in self.STANDARD_LIBS and
-                    package not in self.EXCLUDED_PACKAGES and
-                    not package.startswith('_')):
+            package = node.module.split(".")[0]
+            if (
+                package not in self.STANDARD_LIBS
+                and package not in self.EXCLUDED_PACKAGES
+                and not package.startswith("_")
+            ):
                 self.imports.add(package)
         self.generic_visit(node)
 
@@ -188,23 +231,21 @@ def scan_file_for_imports(file_path: str) -> Set[str]:
         return visitor.imports
 
 
-def write_requirements_file(file_path: str, output_dir: str = None) -> str:
+def write_requirements_file(file_path: str) -> str:
     """
     Scan a Python file for imports and write them to requirements.txt.
 
     Args:
         file_path: Path to the Python file to scan
-        output_dir: Directory where requirements.txt should be created (defaults to parent directory)
 
     Returns:
         Path to the generated requirements.txt file
     """
     imports = scan_file_for_imports(file_path)
 
-    if not output_dir:
-        # Use the parent directory rather than same directory as the file
-        file_dir = os.path.dirname(file_path)
-        output_dir = os.path.dirname(file_dir) if file_dir else "."
+    # Use the parent directory rather than same directory as the file
+    file_dir = os.path.dirname(file_path)
+    output_dir = os.path.dirname(file_dir) if file_dir else "."
 
     requirements_path = os.path.join(output_dir, "requirements.txt")