Format source files

rvermeulen · rvermeulen · commit c6f75d8d54a5 · 2023-01-17T13:48:27.000-08:00
diff --git a/codeql_postproc/cli.py b/codeql_postproc/cli.py
@@ -3,34 +3,53 @@
 from sys import exit
 from typing import Optional, cast
 
+
 @click.group()
 def cli() -> None:
     pass
 
+
 @cli.group()
 def database() -> None:
     pass
 
+
 @database.command("add-vcs-provenance")
-@click.option("-u", "--repository-uri", required=True, help="An absolute URI that specifies the location of the repository.")
-@click.option("-r", "--revision-id", required=True, help="A string that uniquely and permanently identifies the revision.")
+@click.option(
+    "-u",
+    "--repository-uri",
+    required=True,
+    help="An absolute URI that specifies the location of the repository.",
+)
+@click.option(
+    "-r",
+    "--revision-id",
+    required=True,
+    help="A string that uniquely and permanently identifies the revision.",
+)
 @click.argument("database", type=click.Path(exists=True, path_type=Path), required=True)
-def database_add_provenance(repository_uri: str, revision_id: str, database: Path) -> None:
+def database_add_provenance(
+    repository_uri: str, revision_id: str, database: Path
+) -> None:
     from codeql_postproc.helpers.codeql import CodeQLDatabase, InvalidCodeQLDatabase
 
     try:
         codeql_db = CodeQLDatabase(database)
-        vcs_provenance = [{
-            'repositoryUri': repository_uri,
-            'revisionId': revision_id
-        }]
+        vcs_provenance = [{"repositoryUri": repository_uri, "revisionId": revision_id}]
         codeql_db.set_property(versionControlProvenance=vcs_provenance)
     except InvalidCodeQLDatabase as e:
         click.echo(e, err=True)
         exit(1)
 
+
 @database.command("get-property")
-@click.option("-f", "--format", "output_format", type=click.Choice(["json", "yaml"]), default="json")
+@click.option(
+    "-f",
+    "--format",
+    "output_format",
+    type=click.Choice(["json", "yaml"]),
+    default="json",
+)
 @click.argument("key", required=True)
 @click.argument("database", type=click.Path(exists=True, path_type=Path), required=True)
 def get_property(output_format: str, key: str, database: Path) -> None:
@@ -42,57 +61,99 @@ def get_property(output_format: str, key: str, database: Path) -> None:
     try:
         codeql_db = CodeQLDatabase(database)
         value = codeql_db.get_property(key)
-        if output_format == "yaml": 
-            yaml.dump(value, sys.stdout)
-        elif output_format == "json":
-            json.dump(value, sys.stdout)
+        if value:
+            if output_format == "yaml":
+                yaml.dump(value, sys.stdout)
+            elif output_format == "json":
+                json.dump(value, sys.stdout)
+            else:
+                click.echo(f"Unimplemented output format {output_format}!")
+                exit(1)
         else:
-            click.echo(f"Unimplemented output format {output_format}!")
+            click.echo(
+                f"The database does not have a property with key {key}.", err=True
+            )
             exit(1)
-    except KeyError:
-        click.echo(f"The database does not have a property with key {key}.", err=True)
-        exit(1)
     except InvalidCodeQLDatabase as e:
         click.echo(e, err=True)
         exit(1)
 
+
 @cli.group("sarif")
 def sarif() -> None:
     pass
 
+
 @sarif.command("add-vcs-provenance")
 @click.option("-d", "--from-database", is_flag=True)
-@click.option("-u", "--repository-uri", help="An absolute URI that specifies the location of the repository.")
-@click.option("-r", "--revision-id", help="A string that uniquely and permanently identifies the revision.")
-@click.argument("sarif_path", type=click.Path(exists=True, path_type=Path, dir_okay=False), required=True)
-@click.argument("database_path", type=click.Path(exists=True, path_type=Path), required=False)
-def sarif_add_provenance(from_database: bool, repository_uri: str, revision_id: str, sarif_path: Path, database_path: Optional[Path]) -> None:
+@click.option(
+    "-u",
+    "--repository-uri",
+    help="An absolute URI that specifies the location of the repository.",
+)
+@click.option(
+    "-r",
+    "--revision-id",
+    help="A string that uniquely and permanently identifies the revision.",
+)
+@click.argument(
+    "sarif_path",
+    type=click.Path(exists=True, path_type=Path, dir_okay=False),
+    required=True,
+)
+@click.argument(
+    "database_path", type=click.Path(exists=True, path_type=Path), required=False
+)
+def sarif_add_provenance(
+    from_database: bool,
+    repository_uri: str,
+    revision_id: str,
+    sarif_path: Path,
+    database_path: Optional[Path],
+) -> None:
     from codeql_postproc.helpers.codeql import CodeQLDatabase, InvalidCodeQLDatabase
     from codeql_postproc.helpers.sarif import Sarif, InvalidSarif
 
     if from_database and not database_path:
-        raise click.BadArgumentUsage("A database must be specified when using the --from-database option!")
-    
+        raise click.BadArgumentUsage(
+            "A database must be specified when using the --from-database option!"
+        )
+
     if not from_database and not repository_uri:
-        raise click.BadOptionUsage("--repository-uri", "The option '--repository-uri' must be specified if not importing from a database!")
+        raise click.BadOptionUsage(
+            "--repository-uri",
+            "The option '--repository-uri' must be specified if not importing from a database!",
+        )
     if not from_database and not revision_id:
-        raise click.BadOptionUsage("--revision-id", "The option '--revision-id' must be specified if not importing from a database!")
+        raise click.BadOptionUsage(
+            "--revision-id",
+            "The option '--revision-id' must be specified if not importing from a database!",
+        )
 
     if from_database:
         try:
             codeql_db = CodeQLDatabase(cast(Path, database_path))
             # Assume there is only one array element for now.
             vcp = codeql_db.get_property("versionControlProvenance")[0]
             if not "repositoryUri" in vcp:
-                click.echo(f"The database's version control provenance misses the 'repositoryUri' property!", err=True)
+                click.echo(
+                    f"The database's version control provenance misses the 'repositoryUri' property!",
+                    err=True,
+                )
                 exit(1)
             repository_uri = vcp["repositoryUri"]
             if not "revisionId" in vcp:
-                click.echo(f"The database's version control provenance misses the 'revisionId' property!", err=True)
+                click.echo(
+                    f"The database's version control provenance misses the 'revisionId' property!",
+                    err=True,
+                )
                 exit(1)
             revision_id = vcp["revisionId"]
         except KeyError:
-            click.echo(f"The database does not have any version control provenance property.", err=True)
+            click.echo(
+                f"The database does not have any version control provenance property.",
+                err=True,
+            )
             exit(1)
         except InvalidCodeQLDatabase as e:
             click.echo(e, err=True)
@@ -105,7 +166,6 @@ def sarif_add_provenance(from_database: bool, repository_uri: str, revision_id:
         click.echo(f"Unable to process invalid Sarif file with reason: {e}")
         exit(1)
 
-   
-        
+
 if __name__ == "__main__":
-    cli()
+    cli()
diff --git a/codeql_postproc/helpers/codeql.py b/codeql_postproc/helpers/codeql.py
@@ -6,9 +6,11 @@
 from tempfile import TemporaryDirectory
 from jsonpointer import JsonPointer
 
+
 class InvalidCodeQLDatabase(Exception):
     pass
 
+
 class CodeQLDatabase:
     def __init__(self, database: Path) -> None:
         if not database.exists():
@@ -20,24 +22,35 @@ def __init__(self, database: Path) -> None:
                 with db_metadata.open() as fd:
                     database_info = yaml.safe_load(fd)
             else:
-                raise InvalidCodeQLDatabase("Invalid database, missing 'codeql-database.yml'!")
+                raise InvalidCodeQLDatabase(
+                    "Invalid database, missing 'codeql-database.yml'!"
+                )
         elif is_zipfile(database):
             with ZipFile(database) as fd:
+
                 def is_db_metadata(zi: ZipInfo) -> bool:
                     path_parts = zi.filename.split("/")
-                    return len(path_parts) == 2 and path_parts[1] == "codeql-database.yml"
+                    return (
+                        len(path_parts) == 2 and path_parts[1] == "codeql-database.yml"
+                    )
 
                 db_metadata_candidates = list(filter(is_db_metadata, fd.infolist()))
                 if len(db_metadata_candidates) == 0:
-                    raise InvalidCodeQLDatabase("Invalid database, missing 'codeql-database.yml'!")
+                    raise InvalidCodeQLDatabase(
+                        "Invalid database, missing 'codeql-database.yml'!"
+                    )
                 elif len(db_metadata_candidates) > 1:
-                    raise InvalidCodeQLDatabase("Invalid database, found multiple 'codeql-database.yml'!")
+                    raise InvalidCodeQLDatabase(
+                        "Invalid database, found multiple 'codeql-database.yml'!"
+                    )
                 else:
                     self.db_metadata_filename = db_metadata_candidates[0].filename
                     database_info = yaml.safe_load(fd.read(self.db_metadata_filename))
-                
+
         else:
-            raise InvalidCodeQLDatabase("Expected a database directory or database zip archive!")
+            raise InvalidCodeQLDatabase(
+                "Expected a database directory or database zip archive!"
+            )
 
         self.database_info = database_info
         self.database = database
@@ -54,19 +67,21 @@ def update_or_set_props(database: Path, **kwargs: Any):
                 with user_properties.open("r") as fd:
                     existing_props = yaml.safe_load(fd)
                     if existing_props and not isinstance(existing_props, dict):
-                        raise InvalidCodeQLDatabase("The 'user-properties.yml' is not a YAML dictionary!")
+                        raise InvalidCodeQLDatabase(
+                            "The 'user-properties.yml' is not a YAML dictionary!"
+                        )
                     else:
-                        props =  existing_props | props
+                        props = existing_props | props
             with user_properties.open("w") as fd:
                 fd.write(yaml.dump(props))
 
         if self.database.is_dir():
-           update_or_set_props(self.database, **kwargs)
+            update_or_set_props(self.database, **kwargs)
         else:
             with TemporaryDirectory() as tmp_dir:
                 with ZipFile(str(self.database), mode="r") as fd:
                     fd.extractall(tmp_dir)
-                database = Path(tmp_dir) / self.db_metadata_filename.split('/')[0]
+                database = Path(tmp_dir) / self.db_metadata_filename.split("/")[0]
                 update_or_set_props(database, **kwargs)
                 with ZipFile(self.database, mode="w") as fd:
                     for f in database.glob("**/*"):
@@ -79,11 +94,12 @@ def __translate_key(self, key: str) -> JsonPointer:
         if not path.startswith("/"):
             path = "/" + path
         return JsonPointer(path)
-        
 
     def get_property(self, key: str) -> Any:
 
-        database_property = self.__translate_key(key).get(self.database_info, default=None)
+        database_property = self.__translate_key(key).get(
+            self.database_info, default=None
+        )
         if database_property is not None:
             return database_property
 
@@ -94,14 +110,19 @@ def get_property(self, key: str) -> Any:
                 with user_properties_path.open() as fd:
                     user_properties = yaml.safe_load(fd)
         else:
+
             def is_user_property_file(zi: ZipInfo) -> bool:
                 return zi.filename.endswith("user-properties.yml")
+
             with ZipFile(str(self.database)) as fd:
-                user_property_candidates = list(filter(is_user_property_file, fd.infolist()))
+                user_property_candidates = list(
+                    filter(is_user_property_file, fd.infolist())
+                )
                 if len(user_property_candidates) == 0:
                     raise InvalidCodeQLDatabase("Found no 'user-properties.yml' files!")
                 if len(user_property_candidates) > 1:
-                    raise InvalidCodeQLDatabase("Found multiple 'user-properties.yml' files!")
+                    raise InvalidCodeQLDatabase(
+                        "Found multiple 'user-properties.yml' files!"
+                    )
                 user_properties = yaml.safe_load(fd.read(user_property_candidates[0]))
         return self.__translate_key(key).get(user_properties, default=None)
-    
diff --git a/codeql_postproc/helpers/sarif.py b/codeql_postproc/helpers/sarif.py
@@ -4,43 +4,51 @@
 import jsonschema
 import jsonschema.exceptions
 
+
 class InvalidSarif(Exception):
     pass
 
+
 class Sarif:
     def __init__(self, path: Path) -> None:
         self.path = path
 
-        schema_path = Path(__file__).parent.parent / "schemas" / "sarif-schema-2.1.0.json"
+        schema_path = (
+            Path(__file__).parent.parent / "schemas" / "sarif-schema-2.1.0.json"
+        )
         self.schema = json.loads(schema_path.read_text())
-      
+
         try:
             self.content: Dict[str, Any] = json.loads(self.path.read_text())
             jsonschema.validate(self.content, self.schema)
         except json.decoder.JSONDecodeError:
             raise InvalidSarif("Invalid JSON file!")
 
-    def add_version_control_provenance(self, repository_url: str, revision_id: str) -> None:
+    def add_version_control_provenance(
+        self, repository_url: str, revision_id: str
+    ) -> None:
         if not "runs" in self.content or len(self.content["runs"]) == 0:
             raise InvalidSarif("Missing or no run objects in 'runs' property!")
 
         for run in self.content["runs"]:
             if "versionControlProvenance" in run:
                 vcp: Any = run["versionControlProvenance"]
                 if not isinstance(vcp, list):
-                    raise InvalidSarif("The 'versionControlProvenance' property is not an array!")
+                    raise InvalidSarif(
+                        "The 'versionControlProvenance' property is not an array!"
+                    )
             else:
                 run["versionControlProvenance"] = []
-        
-            cast(List[Dict[str, str]], run["versionControlProvenance"]).append({
-                "repositoryUri": repository_url,
-                "revisionId": revision_id
-            })
+
+            cast(List[Dict[str, str]], run["versionControlProvenance"]).append(
+                {"repositoryUri": repository_url, "revisionId": revision_id}
+            )
 
         try:
             jsonschema.validate(self.content, self.schema)
         except jsonschema.exceptions.ValidationError as e:
-            raise InvalidSarif(f"Adding the version control provenance information results in an invalid Sarif file because {e.message}!")
+            raise InvalidSarif(
+                f"Adding the version control provenance information results in an invalid Sarif file because {e.message}!"
+            )
         with self.path.open(mode="w") as fd:
             json.dump(self.content, fd)
-
