Skip to content

Commit 0d7cbfe

Browse files
Timon Violatimonviola
authored andcommitted
chore: apply ruff format
1 parent 8ce10da commit 0d7cbfe

File tree

13 files changed

+19
-13
lines changed

13 files changed

+19
-13
lines changed

src/dagcellent/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""Common Airflow dags and plugins for DFDS."""
2+
23
# SPDX-FileCopyrightText: 2024-present Timon Viola <viotimo@dfds.com>
34
#
45
# SPDX-License-Identifier: MIT

src/dagcellent/_connection.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""Airflow connection utilities."""
2+
23
from __future__ import annotations
34

45
from enum import Enum

src/dagcellent/dag/_meta.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""DAG metadata related utilities."""
2+
23
from __future__ import annotations
34

45
import os

src/dagcellent/data_utils/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""Utilities related to handling various data sources e.g.: databases, cloud blob storage."""
2+
23
from __future__ import annotations
34

45
# pyright: reportUnknownVariableType=false

src/dagcellent/data_utils/sql_reflection.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""SqlAlchemy based utility functions to interact with databases."""
2+
23
from __future__ import annotations
34

45
import logging

src/dagcellent/operators/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""Reusable operators to build Airflow DAGs."""
2+
23
from __future__ import annotations
34

45
from dagcellent.operators.external_table_arrow import CreateExternalTableArrow

src/dagcellent/operators/mlflow/_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""Mlflow helpers."""
2+
23
# ruff: noqa: G004
34
# typing: ignore
45
from __future__ import annotations
@@ -7,7 +8,6 @@
78
from typing import TYPE_CHECKING, TypedDict
89

910
if TYPE_CHECKING:
10-
1111
import mlflow.entities.model_registry
1212

1313

src/dagcellent/operators/mlflow/operators.py

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
from __future__ import annotations
32

43
import logging
@@ -52,7 +51,7 @@ def __init__(
5251
self.version = version
5352
self.tag = tag
5453

55-
super().__init__(**kwargs) # type: ignore [reportUnknownMemberType]
54+
super().__init__(**kwargs) # type: ignore [reportUnknownMemberType]
5655

5756
def execute(self, context: Any) -> None:
5857
"""Operator execute method.
@@ -95,8 +94,7 @@ def __init__(
9594
self.stage = stage
9695
self.model_name = model_name
9796

98-
super().__init__(**kwargs) # type: ignore [reportUnknownMemberType]
99-
97+
super().__init__(**kwargs) # type: ignore [reportUnknownMemberType]
10098

10199
def execute(self, context: Any) -> list[SlimModelVersion]:
102100
"""Connect to Mlflow hook and dump messages to tmp file."""
@@ -131,7 +129,7 @@ def __init__(
131129
self.tracking_uri = tracking_uri
132130
self.upstream_task_id = upstream_task_id
133131

134-
super().__init__(**kwargs)# type: ignore [reportUnknownMemberType]
132+
super().__init__(**kwargs) # type: ignore [reportUnknownMemberType]
135133

136134
def execute(self, context: Any) -> Any:
137135
"""Operator execute method.
@@ -147,13 +145,13 @@ def execute(self, context: Any) -> Any:
147145
key="return_value", task_ids=self.upstream_task_id
148146
)
149147
if isinstance(xcom_val, list):
150-
_first_item = xcom_val.pop() # type: ignore [unknownTypeIssue]
151-
run_id = _first_item["run_id"]# type: ignore [unknownTypeIssue]
148+
_first_item = xcom_val.pop() # type: ignore [unknownTypeIssue]
149+
run_id = _first_item["run_id"] # type: ignore [unknownTypeIssue]
152150
logging.warning(
153151
"Multiple items returned from upstream task. Defaulting to first item"
154152
)
155153
elif isinstance(xcom_val, dict):
156-
run_id = xcom_val["run_id"]# type: ignore [unknownTypeIssue]
154+
run_id = xcom_val["run_id"] # type: ignore [unknownTypeIssue]
157155
else:
158156
logging.error(
159157
"Operator has no implementation for this type: %s", f"{type(xcom_val)}."
@@ -164,8 +162,7 @@ def execute(self, context: Any) -> Any:
164162
raise ValueError(f"{run_id} must be an instance of 'str'")
165163
data = client_hook.get_run(run_id)
166164
logging.info("%s", f"{data=}")
167-
return data.data.metrics # type: ignore [unknownTypeIssue]
168-
165+
return data.data.metrics # type: ignore [unknownTypeIssue]
169166

170167

171168
class GetLatestModelVersion(BaseOperator):
@@ -192,7 +189,7 @@ def __init__(
192189
self.tracking_uri = tracking_uri
193190
self.model_name = model_name
194191

195-
super().__init__(**kwargs)# type: ignore [reportUnknownMemberType]
192+
super().__init__(**kwargs) # type: ignore [reportUnknownMemberType]
196193

197194
def execute(self, context: Any) -> SlimModelVersion:
198195
"""Operator execute method.

tests/dags/test_msql_reflect_otherdb.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
schedule="@once",
1717
catchup=False,
1818
) as dag:
19-
2019
reflect_table = SQLReflectOperator(
2120
task_id="reflect_database", conn_id=CONN_ID, table_name="kaka", database="model"
2221
)

tests/dags/test_mssql_reflect.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Test DAG to show the usage of SQL reflection
22
and executing the returned query.
33
"""
4+
45
from __future__ import annotations
56

67
import datetime

0 commit comments

Comments
 (0)