openlineage-python 1.30.1__tar.gz → 1.32.0__tar.gz

{openlineage_python-1.30.1 → openlineage_python-1.32.0}/PKG-INFO

@@ -1,28 +1,32 @@
 Metadata-Version: 2.4
 Name: openlineage-python
-Version: 1.30.1
+Version: 1.32.0
 Summary: OpenLineage Python Client
 Author-email: OpenLineage <info@openlineage.io>
 Keywords: openlineage
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Python: >=3.8
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.9
 Requires-Dist: attrs>=20.0
 Requires-Dist: packaging>=21.0
 Requires-Dist: python-dateutil>=2.8.2
 Requires-Dist: pyyaml>=5.4
-Requires-Dist: requests>=2.20.0
+Requires-Dist: requests>=2.32.3
 Provides-Extra: docs
 Requires-Dist: beautifulsoup4>=4.12.0; extra == 'docs'
 Requires-Dist: pydata-sphinx-theme>=0.13; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints>=2.0.0; extra == 'docs'
 Requires-Dist: sphinx>=7.0; extra == 'docs'
+Provides-Extra: generator
+Requires-Dist: click; extra == 'generator'
+Requires-Dist: datamodel-code-generator==0.25.4; extra == 'generator'
+Requires-Dist: ruff; extra == 'generator'
 Provides-Extra: kafka
 Requires-Dist: confluent-kafka>=2.1.1; extra == 'kafka'
 Provides-Extra: msk-iam

{openlineage_python-1.30.1 → openlineage_python-1.32.0}/openlineage/client/client.py

@@ -153,7 +153,7 @@ class OpenLineageClient:
     ) -> Event | None:
         """Filters jobs according to config-defined events"""
         for _filter in self._filters:
-            if isinstance(event, RunEvent) and _filter.filter_event(event) is None:
+            if isinstance(event, (RunEvent, event_v2.RunEvent)) and _filter.filter_event(event) is None:
                 return None
         return event
 
@@ -163,6 +163,11 @@ class OpenLineageClient:
         if type(event) not in get_args(Event):
             msg = "`emit` only accepts RunEvent, DatasetEvent, JobEvent classes"
             raise ValueError(msg)
+
+        if log.isEnabledFor(logging.DEBUG):
+            val = Serde.to_json(event).encode("utf-8")
+            log.debug("OpenLineageClient will *try* to emit event %s", val)
+
         if not self.transport:
             log.error("Tried to emit OpenLineage event, but transport is not configured.")
             return
@@ -176,9 +181,6 @@ class OpenLineageClient:
         event = self.add_environment_facets(event)
         event = self.update_event_tags_facets(event)
 
-        if log.isEnabledFor(logging.DEBUG):
-            val = Serde.to_json(event).encode("utf-8")
-            log.debug("OpenLineageClient will emit event %s", val)
         self.transport.emit(event)
         log.debug("OpenLineage event successfully emitted.")
 
@@ -395,7 +397,10 @@ class OpenLineageClient:
         """
         Adds environment variables as facets to the event object.
         """
-        if isinstance(event, RunEvent) and (env_vars := self._collect_environment_variables()):
+        if isinstance(event, (RunEvent, event_v2.RunEvent)) and (
+            env_vars := self._collect_environment_variables()
+        ):
+            event.run.facets = event.run.facets or {}
             event.run.facets["environmentVariables"] = EnvironmentVariablesRunFacet(
                 environmentVariables=[
                     EnvironmentVariable(name=name, value=value) for name, value in env_vars.items()

{openlineage_python-1.30.1 → openlineage_python-1.32.0}/openlineage/client/constants.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-__version__ = "1.30.1"
+__version__ = "1.32.0"
 
 DEFAULT_TIMEOUT_MS = 5000
 DEFAULT_NAMESPACE_NAME = "default"

{openlineage_python-1.30.1 → openlineage_python-1.32.0}/openlineage/client/facet.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import warnings
 from enum import Enum
-from typing import Any, ClassVar, Dict, List, Optional
+from typing import Any, ClassVar, Optional
 
 import attr
 from openlineage.client.constants import DEFAULT_PRODUCER
@@ -27,8 +27,8 @@ class BaseFacet(RedactMixin):
     _producer: str = attr.ib(init=False)
     _schemaURL: str = attr.ib(init=False)  # noqa: N815
 
-    _base_skip_redact: ClassVar[List[str]] = ["_producer", "_schemaURL"]
-    _additional_skip_redact: ClassVar[List[str]] = []
+    _base_skip_redact: ClassVar[list[str]] = ["_producer", "_schemaURL"]
+    _additional_skip_redact: ClassVar[list[str]] = []
 
     def __attrs_post_init__(self) -> None:
         self._producer = PRODUCER
@@ -39,7 +39,7 @@ class BaseFacet(RedactMixin):
         return SCHEMA_URI + "#/definitions/BaseFacet"
 
     @property
-    def skip_redact(self) -> List[str]:
+    def skip_redact(self) -> list[str]:
         return self._base_skip_redact + self._additional_skip_redact
 
 
@@ -48,7 +48,7 @@ class NominalTimeRunFacet(BaseFacet):
     nominalStartTime: str = attr.ib()  # noqa: N815
     nominalEndTime: Optional[str] = attr.ib(default=None)  # noqa: N815
 
-    _additional_skip_redact: ClassVar[List[str]] = ["nominalStartTime", "nominalEndTime"]
+    _additional_skip_redact: ClassVar[list[str]] = ["nominalStartTime", "nominalEndTime"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -57,10 +57,10 @@ class NominalTimeRunFacet(BaseFacet):
 
 @attr.s
 class ParentRunFacet(BaseFacet):
-    run: Dict[Any, Any] = attr.ib()
-    job: Dict[Any, Any] = attr.ib()
+    run: dict[Any, Any] = attr.ib()
+    job: dict[Any, Any] = attr.ib()
 
-    _additional_skip_redact: ClassVar[List[str]] = ["job", "run"]
+    _additional_skip_redact: ClassVar[list[str]] = ["job", "run"]
 
     @classmethod
     def create(cls, runId: str, namespace: str, name: str) -> "ParentRunFacet":  # noqa: N803
@@ -93,7 +93,7 @@ class SourceCodeLocationJobFacet(BaseFacet):
     type: str = attr.ib()
     url: str = attr.ib()
 
-    _additional_skip_redact: ClassVar[List[str]] = ["type", "url"]
+    _additional_skip_redact: ClassVar[list[str]] = ["type", "url"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -124,12 +124,12 @@ class SchemaField(RedactMixin):
     type: str = attr.ib()
     description: Optional[str] = attr.ib(default=None)
 
-    _do_not_redact: ClassVar[List[str]] = ["name", "type"]
+    _do_not_redact: ClassVar[list[str]] = ["name", "type"]
 
 
 @attr.s
 class SchemaDatasetFacet(BaseFacet):
-    fields: List[SchemaField] = attr.ib()
+    fields: list[SchemaField] = attr.ib()
 
     @staticmethod
     def _get_schema() -> str:
@@ -141,7 +141,7 @@ class DataSourceDatasetFacet(BaseFacet):
     name: str = attr.ib()
     uri: str = attr.ib()
 
-    _additional_skip_redact: ClassVar[List[str]] = ["name", "uri"]
+    _additional_skip_redact: ClassVar[list[str]] = ["name", "uri"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -154,7 +154,7 @@ class OutputStatisticsOutputDatasetFacet(BaseFacet):
     size: Optional[int] = attr.ib(default=None)
     fileCount: Optional[int] = attr.ib(default=None)  # noqa: N815
 
-    _additional_skip_redact: ClassVar[List[str]] = ["rowCount", "size", "fileCount"]
+    _additional_skip_redact: ClassVar[list[str]] = ["rowCount", "size", "fileCount"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -169,7 +169,7 @@ class ColumnMetric:
     count: Optional[int] = attr.ib(default=None)
     min: Optional[float] = attr.ib(default=None)
     max: Optional[float] = attr.ib(default=None)
-    quantiles: Optional[Dict[str, float]] = attr.ib(default=None)
+    quantiles: Optional[dict[str, float]] = attr.ib(default=None)
 
 
 @attr.s
@@ -177,7 +177,7 @@ class DataQualityMetricsInputDatasetFacet(BaseFacet):
     rowCount: Optional[int] = attr.ib(default=None)  # noqa: N815
     bytes: Optional[int] = attr.ib(default=None)
     fileCount: Optional[int] = attr.ib(default=None)  # noqa: N815
-    columnMetrics: Dict[str, ColumnMetric] = attr.ib(factory=dict)  # noqa: N815
+    columnMetrics: dict[str, ColumnMetric] = attr.ib(factory=dict)  # noqa: N815
 
     @staticmethod
     def _get_schema() -> str:
@@ -190,7 +190,7 @@ class Assertion(RedactMixin):
     success: bool = attr.ib()
     column: Optional[str] = attr.ib(default=None)
 
-    _skip_redact: ClassVar[List[str]] = ["column"]
+    _skip_redact: ClassVar[list[str]] = ["column"]
 
 
 @attr.s
@@ -198,7 +198,7 @@ class DataQualityAssertionsDatasetFacet(BaseFacet):
 
     """This facet represents asserted expectations on dataset or it's column."""
 
-    assertions: List[Assertion] = attr.ib()
+    assertions: list[Assertion] = attr.ib()
 
     @staticmethod
     def _get_schema() -> str:
@@ -213,7 +213,7 @@ class SourceCodeJobFacet(BaseFacet):
     language: str = attr.ib()  # language that the code was written in
     source: str = attr.ib()  # source code text
 
-    _additional_skip_redact: ClassVar[List[str]] = ["language"]
+    _additional_skip_redact: ClassVar[list[str]] = ["language"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -235,7 +235,7 @@ class ErrorMessageRunFacet(BaseFacet):
     programmingLanguage: str = attr.ib()  # noqa: N815
     stackTrace: Optional[str] = attr.ib(default=None)  # noqa: N815
 
-    _additional_skip_redact: ClassVar[List[str]] = ["programmingLanguage"]
+    _additional_skip_redact: ClassVar[list[str]] = ["programmingLanguage"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -254,7 +254,7 @@ class SymlinksDatasetFacet(BaseFacet):
 
     """This facet represents dataset symlink names."""
 
-    identifiers: List[SymlinksDatasetFacetIdentifiers] = attr.ib(factory=dict)
+    identifiers: list[SymlinksDatasetFacetIdentifiers] = attr.ib(factory=dict)
 
     @staticmethod
     def _get_schema() -> str:
@@ -285,7 +285,7 @@ class OwnershipJobFacet(BaseFacet):
 
     """This facet represents ownership of a job."""
 
-    owners: List[OwnershipJobFacetOwners] = attr.ib(factory=dict)
+    owners: list[OwnershipJobFacetOwners] = attr.ib(factory=dict)
 
     @staticmethod
     def _get_schema() -> str:
@@ -357,7 +357,7 @@ class OwnershipDatasetFacet(BaseFacet):
 
     """This facet represents ownership of a dataset."""
 
-    owners: List[OwnershipDatasetFacetOwners] = attr.ib(factory=dict)
+    owners: list[OwnershipDatasetFacetOwners] = attr.ib(factory=dict)
 
     @staticmethod
     def _get_schema() -> str:
@@ -370,12 +370,12 @@ class ColumnLineageDatasetFacetFieldsAdditionalInputFields(RedactMixin):
     name: str = attr.ib()
     field: str = attr.ib()
 
-    _skip_redact: ClassVar[List[str]] = ["namespace", "name", "field"]
+    _skip_redact: ClassVar[list[str]] = ["namespace", "name", "field"]
 
 
 @attr.s
 class ColumnLineageDatasetFacetFieldsAdditional:
-    inputFields: ClassVar[List[ColumnLineageDatasetFacetFieldsAdditionalInputFields]] = attr.ib()  # noqa: N815
+    inputFields: ClassVar[list[ColumnLineageDatasetFacetFieldsAdditionalInputFields]] = attr.ib()  # noqa: N815
     transformationDescription: str = attr.ib()  # noqa:  N815
     transformationType: str = attr.ib()  # noqa:  N815
 
@@ -385,7 +385,7 @@ class ColumnLineageDatasetFacet(BaseFacet):
 
     """This facet contains column lineage of a dataset."""
 
-    fields: Dict[str, ColumnLineageDatasetFacetFieldsAdditional] = attr.ib(factory=dict)
+    fields: dict[str, ColumnLineageDatasetFacetFieldsAdditional] = attr.ib(factory=dict)
 
     @staticmethod
     def _get_schema() -> str:
@@ -415,7 +415,7 @@ class ExtractionError(BaseFacet):
 class ExtractionErrorRunFacet(BaseFacet):
     totalTasks: int = attr.ib()  # noqa: N815
     failedTasks: int = attr.ib()  # noqa: N815
-    errors: List[ExtractionError] = attr.ib()
+    errors: list[ExtractionError] = attr.ib()
 
     @staticmethod
     def _get_schema() -> str:

{openlineage_python-1.30.1 → openlineage_python-1.32.0}/openlineage/client/filter.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import logging
 import re
 import typing
 
@@ -9,6 +10,7 @@ import attr
 from openlineage.client.event_v2 import RunEvent as RunEvent_v2
 from openlineage.client.run import RunEvent
 
+log = logging.getLogger(__name__)
 RunEventType = typing.Union[RunEvent, RunEvent_v2]
 
 
@@ -30,6 +32,7 @@ class ExactMatchFilter(Filter):
 
     def filter_event(self, event: RunEventType) -> RunEventType | None:
         if self.match == event.job.name:
+            log.debug("Job name `%s` matches exactly `%s`.", event.job.name, self.match)
             return None
         return event
 
@@ -40,16 +43,20 @@ class RegexFilter(Filter):
 
     def filter_event(self, event: RunEventType) -> RunEventType | None:
         if self.pattern.match(event.job.name):
+            log.debug("Job name `%s` matches regex `%s`.", event.job.name, self.pattern.pattern)
             return None
         return event
 
 
 def create_filter(conf: FilterConfig) -> Filter | None:
     if not conf.type:
+        log.warning("OpenLineage filter config must have a `type`.")
         return None
-    # Switch in 3.10 🙂
     if conf.type == "exact" and conf.match:
+        log.debug("Creating ExactMatchFilter with match='%s'", conf.match)
         return ExactMatchFilter(match=conf.match)
     if conf.type == "regex" and conf.regex:
+        log.debug("Creating RegexFilter with regex='%s'", conf.regex)
         return RegexFilter(regex=conf.regex)
+    log.warning("Unsupported OpenLineage filter type: `%s`", conf.type)
     return None

{openlineage_python-1.30.1 → openlineage_python-1.32.0}/openlineage/client/generated/parent_run.py

@@ -28,11 +28,12 @@ class ParentRunFacet(RunFacet):
 
     run: Run
     job: Job
+    root: Root | None = attr.field(default=None)
     _additional_skip_redact: ClassVar[list[str]] = ["job", "run"]
 
     @staticmethod
     def _get_schema() -> str:
-        return "https://openlineage.io/spec/facets/1-0-1/ParentRunFacet.json#/$defs/ParentRunFacet"
+        return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/ParentRunFacet"
 
     @classmethod
     def create(cls, runId: str, namespace: str, name: str) -> ParentRunFacet:  # noqa: N803
@@ -46,6 +47,46 @@ class ParentRunFacet(RunFacet):
         return cls(run=Run(runId=runId), job=Job(namespace=namespace, name=name))
 
 
+@attr.define
+class Root(RedactMixin):
+    run: RootRun
+    job: RootJob
+    _skip_redact: ClassVar[list[str]] = ["run", "job"]
+
+
+@attr.define
+class RootJob(RedactMixin):
+    namespace: str
+    """The namespace containing root job"""
+
+    name: str
+    """The unique name containing root job within that namespace"""
+
+    _skip_redact: ClassVar[list[str]] = ["namespace", "name"]
+
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/RootJob"
+
+
+@attr.define
+class RootRun(RedactMixin):
+    runId: str = attr.field()  # noqa: N815
+    """The globally unique ID of the root run associated with the root job."""
+
+    _skip_redact: ClassVar[list[str]] = ["runId"]
+
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/RootRun"
+
+    @runId.validator
+    def runid_check(self, attribute: str, value: str) -> None:  # noqa: ARG002
+        from uuid import UUID
+
+        UUID(value)
+
+
 @attr.define
 class Run(RedactMixin):
     runId: str = attr.field()  # noqa: N815

openlineage_python-1.32.0/openlineage/client/generator/base.py

@@ -0,0 +1,257 @@
+# Copyright 2018-2025 contributors to the OpenLineage project
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import copy
+import json
+import logging
+import os
+import pathlib
+import re
+import subprocess
+import tempfile
+from typing import Any
+
+from datamodel_code_generator import DataModelType, PythonVersion
+from datamodel_code_generator.imports import Import
+from datamodel_code_generator.model import get_data_model_types
+from datamodel_code_generator.model import pydantic as pydantic_model
+from datamodel_code_generator.parser.jsonschema import JsonSchemaParser
+from datamodel_code_generator.types import Types
+
+log = logging.getLogger(__name__)
+
+
+def camel_to_snake(string: str) -> str:
+    """Convert camel_case string to SnakeCase"""
+    _under_score_1 = re.compile(r"([^_])([A-Z][a-z]+)")
+    _under_score_2 = re.compile("([a-z0-9])([A-Z])")
+    subbed = _under_score_1.sub(r"\1_\2", string)
+    return _under_score_2.sub(r"\1_\2", subbed).lower()
+
+
+# import attrs instead of dataclass
+DATA_MODELS_TYPES = get_data_model_types(
+    DataModelType.DataclassesDataclass, target_python_version=PythonVersion.PY_38
+)
+NEW_MODEL = DATA_MODELS_TYPES.data_model
+NEW_MODEL.DEFAULT_IMPORTS = (Import.from_full_path("attr"),)
+
+# custom code
+SET_PRODUCER_CODE = """
+PRODUCER = DEFAULT_PRODUCER
+
+def set_producer(producer: str) -> None:
+    global PRODUCER  # noqa: PLW0603
+    PRODUCER = producer
+"""
+
+# locations definitions
+FILE_LOCATION = pathlib.Path(__file__).resolve().parent
+TEMPLATES_LOCATION = FILE_LOCATION / "templates"
+
+# Global dictionaries intended to be imported by other modules
+# Contains schema URLs and base IDs parsed from specification files
+SCHEMA_URLS: dict[str, str] = {}
+BASE_IDS: dict[str, str] = {}
+
+
+def deep_merge_dicts(dict1: dict[str, Any], dict2: dict[str, Any]) -> dict[str, Any]:
+    """Deep merges two dictionaries.
+
+    This function merges two dictionaries while handling nested dictionaries.
+    For keys that exist in both dictionaries, the values from dict2 take precedence.
+    If a key exists in both dictionaries and the values are dictionaries themselves,
+    they are merged recursively.
+    This function merges only dictionaries. If key is of different type, e.g. list
+    it does not work properly.
+    """
+    merged = dict1.copy()
+    for k, v in dict2.items():
+        if k in merged and isinstance(v, dict):
+            merged[k] = deep_merge_dicts(merged.get(k, {}), v)
+        else:
+            merged[k] = v
+    return merged
+
+
+def parse_additional_data(spec: dict[str, Any], file_name: str) -> None:
+    """Parse additional data from spec files.
+
+    Parses:
+        * schema URLs
+        * base IDs
+
+    Updates the module-level SCHEMA_URLS and BASE_IDS dictionaries
+    that can be imported by other modules.
+    """
+    base_id = spec["$id"]
+    for name, _ in spec["$defs"].items():
+        SCHEMA_URLS[name] = f"{base_id}#/$defs/{name}"
+        BASE_IDS[file_name] = spec["$id"]
+
+
+def load_specs(base_spec_location: pathlib.Path, facets_spec_location: pathlib.Path) -> list[pathlib.Path]:
+    """Load base `OpenLineage.json` and other facets' spec files"""
+    locations = []
+    if facets_spec_location.is_dir():
+        facets_spec_files = sorted(pathlib.Path(os.path.abspath(facets_spec_location)).glob("*.json"))
+    else:
+        facets_spec_files = [pathlib.Path(os.path.abspath(facets_spec_location))]
+
+    file_specs = [
+        base_spec_location.resolve(),  # Base spec must be the first element
+        *facets_spec_files,
+    ]
+    for file_spec in file_specs:
+        spec = json.loads(file_spec.read_text())
+        parse_additional_data(spec, file_spec.name)
+        locations.append(file_spec)
+
+    return locations
+
+
+def parse_and_generate(
+    locations: list[pathlib.Path], extra_template_data: dict[str, Any] | None = None
+) -> dict[str, Any]:
+    """Parse and generate data models from a given specification."""
+    temporary_locations = []
+
+    current_dir = pathlib.Path.cwd()
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_directory = pathlib.Path(tmp).resolve()
+        for location in locations:
+            tmp_location = (tmp_directory / location.name).resolve()
+            tmp_location.write_text(location.read_text())
+            temporary_locations.append(tmp_location)
+
+        os.chdir(tmp_directory)
+        # first parse OpenLineage.json
+        parser = JsonSchemaParser(
+            source=temporary_locations[:1],
+            data_model_type=NEW_MODEL,
+            data_model_root_type=DATA_MODELS_TYPES.root_model,
+            data_model_field_type=DATA_MODELS_TYPES.field_model,
+            data_type_manager_type=DATA_MODELS_TYPES.data_type_manager,
+            dump_resolve_reference_action=DATA_MODELS_TYPES.dump_resolve_reference_action,
+            special_field_name_prefix="",
+            use_schema_description=True,
+            field_constraints=True,
+            use_union_operator=True,
+            use_standard_collections=True,
+            base_class="openlineage.client.utils.RedactMixin",
+            class_name="ClassToBeSkipped",
+            use_field_description=True,
+            use_double_quotes=True,
+            keep_model_order=True,
+            custom_template_dir=TEMPLATES_LOCATION,
+            extra_template_data=extra_template_data,  # type: ignore[arg-type]
+            additional_imports=[
+                "typing.ClassVar",
+                "typing.Any",
+                "typing.cast",
+                "openlineage.client.constants.DEFAULT_PRODUCER",
+            ],
+        )
+
+        # keep information about uuid and date-time formats
+        # this is going to be changed back to str type hint in jinja template
+        data_type_manager: pydantic_model.DataTypeManager = parser.data_type_manager  # type: ignore[assignment]
+        uuid_type = copy.deepcopy(data_type_manager.type_map[Types.uuid])
+        uuid_type.type = "uuid"
+        data_type_manager.type_map[Types.uuid] = uuid_type
+
+        date_time_type = copy.deepcopy(data_type_manager.type_map[Types.date_time])
+        date_time_type.type = "date-time"
+        data_type_manager.type_map[Types.date_time] = date_time_type
+
+        uri_type = copy.deepcopy(data_type_manager.type_map[Types.date_time])
+        uri_type.type = "uri"
+        data_type_manager.type_map[Types.uri] = uri_type
+
+        parser.parse(format_=False)
+
+        # parse rest of spec
+        parser.source = temporary_locations[1:]
+
+        # change paths so that parser sees base objects as local
+        parser.model_resolver.references = {
+            k.replace("OpenLineage.json", BASE_IDS["OpenLineage.json"]): v
+            for k, v in parser.model_resolver.references.items()
+        }
+
+        output = parser.parse(format_=False)
+
+    # go back to original directory
+    os.chdir(current_dir)
+
+    return output  # type: ignore[return-value]
+
+
+def separate_imports(code: str) -> tuple[str, str]:
+    """Separates a Python script code (as string) into imports and the rest."""
+    imports_section = []
+    rest_of_code = []
+    in_import = False
+    for line in code.splitlines():
+        if line.startswith(("import ", "from ")):
+            in_import = True
+        elif in_import and line.strip() == "":
+            in_import = False
+        if in_import:
+            imports_section.append(line)
+        else:
+            rest_of_code.append(line)
+    return ("\n".join(imports_section), "\n".join(rest_of_code))
+
+
+def format_and_save_output(
+    output: str,
+    location: pathlib.Path,
+    add_set_producer_code: bool | None = False,
+    header: str | None = None,
+) -> None:
+    """Adjust and format generated file."""
+
+    # save temporary file to the same directory that output
+    # so that ruff receives same rules
+    with tempfile.NamedTemporaryFile(
+        "w", prefix=location.stem.lower(), suffix=".py", dir=location.parent, delete=False
+    ) as tmpfile:
+        if header:
+            tmpfile.write(header + "\n")
+        output = output.replace("from .OpenLineage", "from openlineage.client.generated.base")
+        imports_section, rest_of_code = separate_imports(output)
+        tmpfile.write(imports_section)
+        if add_set_producer_code:
+            tmpfile.write(SET_PRODUCER_CODE)
+        tmpfile.write(rest_of_code)
+        tmpfile.flush()
+
+    # run ruff lint
+    with subprocess.Popen(
+        args=["ruff", "check", tmpfile.name, "--fix"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        close_fds=True,
+    ) as lint_process:
+        if lint_process.returncode:
+            log.warning("Ruff lint failed: %s", lint_process.returncode)
+
+    # run ruff format
+    with subprocess.Popen(
+        args=["ruff", "format", tmpfile.name],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        close_fds=True,
+    ) as format_process:
+        if format_process.returncode:
+            log.warning("Ruff lint failed: %s", format_process.returncode)
+
+    # move file to output location
+    if location.name == "open_lineage.py":
+        location = location.with_name("base.py")
+    os.rename(tmpfile.name, location)
+
+    if lint_process.returncode or format_process.returncode:
+        log.warning("%s failed on ruff.", location)