Convert methods load_from_disk and load_yamlfile_from_disk_and_exit to work with all yaml files

Author: dgarros

infrahub_sdk/ctl/schema.py

@@ -11,13 +11,12 @@
 from infrahub_sdk import InfrahubClient
 from infrahub_sdk.async_typer import AsyncTyper
 from infrahub_sdk.ctl.client import initialize_client
-from infrahub_sdk.ctl.exceptions import FileNotValidError
 from infrahub_sdk.ctl.utils import catch_exception, init_logging
 from infrahub_sdk.queries import SCHEMA_HASH_SYNC_STATUS
-from infrahub_sdk.utils import find_files
 from infrahub_sdk.yaml import SchemaFile
 
 from .parameters import CONFIG_PARAM
+from .utils import load_yamlfile_from_disk_and_exit
 
 app = AsyncTyper()
 console = Console()
@@ -30,45 +29,6 @@ def callback() -> None:
     """
 
 
-def load_schemas_from_disk(schemas: list[Path]) -> list[SchemaFile]:
-    schemas_data: list[SchemaFile] = []
-    for schema in schemas:
-        if schema.is_file():
-            schema_file = SchemaFile(location=schema)
-            schema_file.load_content()
-            schemas_data.append(schema_file)
-        elif schema.is_dir():
-            files = find_files(extension=["yaml", "yml", "json"], directory=schema)
-            for item in files:
-                schema_file = SchemaFile(location=item)
-                schema_file.load_content()
-                schemas_data.append(schema_file)
-        else:
-            raise FileNotValidError(name=schema, message=f"Schema path: {schema} does not exist!")
-
-    return schemas_data
-
-
-def load_schemas_from_disk_and_exit(schemas: list[Path]) -> list[SchemaFile]:
-    has_error = False
-    try:
-        schemas_data = load_schemas_from_disk(schemas=schemas)
-    except FileNotValidError as exc:
-        console.print(f"[red]{exc.message}")
-        raise typer.Exit(1) from exc
-
-    for schema_file in schemas_data:
-        if schema_file.valid and schema_file.content:
-            continue
-        console.print(f"[red]{schema_file.error_message} ({schema_file.location})")
-        has_error = True
-
-    if has_error:
-        raise typer.Exit(1)
-
-    return schemas_data
-
-
 def validate_schema_content_and_exit(client: InfrahubClient, schemas: list[SchemaFile]) -> None:
     has_error: bool = False
     for schema_file in schemas:
@@ -153,7 +113,7 @@ async def load(
 
     init_logging(debug=debug)
 
-    schemas_data = load_schemas_from_disk_and_exit(schemas=schemas)
+    schemas_data = load_yamlfile_from_disk_and_exit(paths=schemas, file_type=SchemaFile, console=console)
     schema_definition = "schema" if len(schemas_data) == 1 else "schemas"
     client = await initialize_client()
     validate_schema_content_and_exit(client=client, schemas=schemas_data)
@@ -203,7 +163,7 @@ async def check(
 
     init_logging(debug=debug)
 
-    schemas_data = load_schemas_from_disk_and_exit(schemas=schemas)
+    schemas_data = load_yamlfile_from_disk_and_exit(paths=schemas, file_type=SchemaFile, console=console)
     client = await initialize_client()
     validate_schema_content_and_exit(client=client, schemas=schemas_data)
 

infrahub_sdk/ctl/utils.py

@@ -3,7 +3,7 @@
 import traceback
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional, TypeVar, Union
 
 import pendulum
 import typer
@@ -14,7 +14,7 @@
 from rich.logging import RichHandler
 from rich.markup import escape
 
-from infrahub_sdk.ctl.exceptions import QueryNotFoundError
+from infrahub_sdk.ctl.exceptions import FileNotValidError, QueryNotFoundError
 from infrahub_sdk.exceptions import (
     AuthenticationError,
     Error,
@@ -25,9 +25,12 @@
     ServerNotResponsiveError,
 )
 from infrahub_sdk.schema import InfrahubRepositoryConfig
+from infrahub_sdk.yaml import YamlFile
 
 from .client import initialize_client_sync
 
+YamlFileVar = TypeVar("YamlFileVar", bound=YamlFile)
+
 
 def init_logging(debug: bool = False) -> None:
     logging.getLogger("infrahub_sdk").setLevel(logging.CRITICAL)
@@ -179,3 +182,25 @@ def get_fixtures_dir() -> Path:
     """Get the directory which stores fixtures that are common to multiple unit/integration tests."""
     here = Path(__file__).resolve().parent
     return here.parent.parent / "tests" / "fixtures"
+
+
+def load_yamlfile_from_disk_and_exit(
+    paths: list[Path], file_type: type[YamlFileVar], console: Console
+) -> list[YamlFileVar]:
+    has_error = False
+    try:
+        data_files = file_type.load_from_disk(paths=paths)
+    except FileNotValidError as exc:
+        console.print(f"[red]{exc.message}")
+        raise typer.Exit(1) from exc
+
+    for data_file in data_files:
+        if data_file.valid and data_file.content:
+            continue
+        console.print(f"[red]{data_file.error_message} ({data_file.location})")
+        has_error = True
+
+    if has_error:
+        raise typer.Exit(1)
+
+    return data_files

infrahub_sdk/yaml.py

@@ -1,17 +1,39 @@
+from enum import Enum
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Self
 
 import yaml
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
+from infrahub_sdk.ctl.exceptions import FileNotValidError
+from infrahub_sdk.utils import find_files
 
-class SchemaFile(BaseModel):
+
+class InfrahubFileApiVersion(str, Enum):
+    V1 = "infrahub.app/v1"
+
+
+class InfrahubFileKind(str, Enum):
+    MENU = "Menu"
+    OBJECT = "Object"
+
+
+class InfrahubFileData(BaseModel):
+    api_version: InfrahubFileApiVersion = Field(InfrahubFileApiVersion.V1, alias="apiVersion")
+    kind: InfrahubFileKind
+    spec: dict
+    metadata: Optional[dict] = Field(default_factory=dict)
+
+
+class LocalFile(BaseModel):
     identifier: Optional[str] = None
     location: Path
     content: Optional[dict] = None
     valid: bool = True
     error_message: Optional[str] = None
 
+
+class YamlFile(LocalFile):
     def load_content(self) -> None:
         try:
             self.content = yaml.safe_load(self.location.read_text())
@@ -23,3 +45,52 @@ def load_content(self) -> None:
         if not self.content:
             self.error_message = "Empty YAML/JSON file"
             self.valid = False
+
+    def validate_content(self) -> None:
+        pass
+
+    @classmethod
+    def load_from_disk(cls, paths: list[Path]) -> list[Self]:
+        yaml_files: list[Self] = []
+        for file_path in paths:
+            if file_path.is_file():
+                yaml_file = cls(location=file_path)
+                yaml_file.load_content()
+                yaml_files.append(yaml_file)
+            elif file_path.is_dir():
+                files = find_files(extension=["yaml", "yml", "json"], directory=file_path)
+                for item in files:
+                    yaml_file = cls(location=item)
+                    yaml_file.load_content()
+                    yaml_files.append(yaml_file)
+            else:
+                raise FileNotValidError(name=str(file_path), message=f"{file_path} does not exist!")
+
+        return yaml_files
+
+
+class InfrahubFile(YamlFile):
+    _data: Optional[InfrahubFileData] = None
+
+    @property
+    def data(self) -> InfrahubFileData:
+        if not self._data:
+            raise ValueError("_data hasn't been initialized yet")
+        return self._data
+
+    @property
+    def version(self) -> InfrahubFileApiVersion:
+        return self.data.api_version
+
+    @property
+    def kind(self) -> InfrahubFileKind:
+        return self.data.kind
+
+    def validate_content(self) -> None:
+        if not self.content:
+            raise ValueError("Content hasn't been loaded yet")
+        self._data = InfrahubFileData(**self.content)
+
+
+class SchemaFile(YamlFile):
+    pass