Skip to content

Commit

Permalink
Add import sorting (isort) to Cosmos (astronomer#866)
Browse files Browse the repository at this point in the history
Previously Cosmos didn't have any import sorting set in the
`pyproject.toml`, and I thought it would be nice to have a uniform
import style and clean up the existing imports.

This PR updates adds ruff isort that will now be fixed by pre-commit.
  • Loading branch information
jbandoro authored and arojasb3 committed Jul 14, 2024
1 parent a8246ac commit 8b3441e
Show file tree
Hide file tree
Showing 65 changed files with 163 additions and 174 deletions.
6 changes: 3 additions & 3 deletions cosmos/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,12 @@
from cosmos.airflow.dag import DbtDag
from cosmos.airflow.task_group import DbtTaskGroup
from cosmos.config import (
ProjectConfig,
ProfileConfig,
ExecutionConfig,
ProfileConfig,
ProjectConfig,
RenderConfig,
)
from cosmos.constants import LoadMode, TestBehavior, ExecutionMode
from cosmos.constants import ExecutionMode, LoadMode, TestBehavior
from cosmos.log import get_logger
from cosmos.operators.lazy_load import MissingPackage
from cosmos.operators.local import (
Expand Down
2 changes: 1 addition & 1 deletion cosmos/airflow/dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

from airflow.models.dag import DAG

from cosmos.converter import airflow_kwargs, specific_kwargs, DbtToAirflowConverter
from cosmos.converter import DbtToAirflowConverter, airflow_kwargs, specific_kwargs


class DbtDag(DAG, DbtToAirflowConverter):
Expand Down
12 changes: 5 additions & 7 deletions cosmos/airflow/graph.py
Original file line number Diff line number Diff line change
@@ -1,26 +1,24 @@
from __future__ import annotations

from typing import Any, Callable
from typing import Any, Callable, Union

from airflow.models import BaseOperator
from airflow.models.dag import DAG
from airflow.utils.task_group import TaskGroup

from cosmos.config import RenderConfig
from cosmos.constants import (
DEFAULT_DBT_RESOURCES,
TESTABLE_DBT_RESOURCES,
DbtResourceType,
ExecutionMode,
TestBehavior,
TestIndirectSelection,
ExecutionMode,
TESTABLE_DBT_RESOURCES,
DEFAULT_DBT_RESOURCES,
)
from cosmos.config import RenderConfig
from cosmos.core.airflow import get_airflow_task as create_airflow_task
from cosmos.core.graph.entities import Task as TaskMetadata
from cosmos.dbt.graph import DbtNode
from cosmos.log import get_logger
from typing import Union


logger = get_logger(__name__)

Expand Down
3 changes: 2 additions & 1 deletion cosmos/airflow/task_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,12 @@
"""

from __future__ import annotations

from typing import Any

from airflow.utils.task_group import TaskGroup

from cosmos.converter import airflow_kwargs, specific_kwargs, DbtToAirflowConverter
from cosmos.converter import DbtToAirflowConverter, airflow_kwargs, specific_kwargs


class DbtTaskGroup(TaskGroup, DbtToAirflowConverter):
Expand Down
8 changes: 4 additions & 4 deletions cosmos/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,18 @@
import contextlib
import shutil
import tempfile
import warnings
from dataclasses import InitVar, dataclass, field
from pathlib import Path
import warnings
from typing import Any, Iterator, Callable
from typing import Any, Callable, Iterator

from cosmos.constants import (
DbtResourceType,
TestBehavior,
ExecutionMode,
InvocationMode,
LoadMode,
TestBehavior,
TestIndirectSelection,
InvocationMode,
)
from cosmos.dbt.executable import get_system_dbt
from cosmos.exceptions import CosmosValueError
Expand Down
1 change: 0 additions & 1 deletion cosmos/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import aenum


DBT_PROFILE_PATH = Path(os.path.expanduser("~")).joinpath(".dbt/profiles.yml")
DEFAULT_DBT_PROFILE_NAME = "cosmos_profile"
DEFAULT_DBT_TARGET_NAME = "cosmos_target"
Expand Down
4 changes: 2 additions & 2 deletions cosmos/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,19 @@

from __future__ import annotations

import copy
import inspect
from typing import Any, Callable
import copy
from warnings import warn

from airflow.models.dag import DAG
from airflow.utils.task_group import TaskGroup

from cosmos.airflow.graph import build_airflow_graph
from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig
from cosmos.constants import ExecutionMode
from cosmos.dbt.graph import DbtGraph
from cosmos.dbt.selector import retrieve_by_label
from cosmos.config import ProjectConfig, ExecutionConfig, RenderConfig, ProfileConfig
from cosmos.exceptions import CosmosValueError
from cosmos.log import get_logger

Expand Down
1 change: 0 additions & 1 deletion cosmos/core/airflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from cosmos.core.graph.entities import Task
from cosmos.log import get_logger


logger = get_logger(__name__)


Expand Down
5 changes: 3 additions & 2 deletions cosmos/dbt/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,13 @@
import json
import os
import tempfile
import yaml
from dataclasses import dataclass, field
from pathlib import Path
from subprocess import PIPE, Popen
from typing import Any

import yaml

from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig
from cosmos.constants import (
DBT_LOG_DIR_NAME,
Expand All @@ -22,7 +23,7 @@
LoadMode,
)
from cosmos.dbt.parser.project import LegacyDbtProject
from cosmos.dbt.project import create_symlinks, copy_msgpack_for_partial_parse, environ
from cosmos.dbt.project import copy_msgpack_for_partial_parse, create_symlinks, environ
from cosmos.dbt.selector import select_nodes
from cosmos.log import get_logger

Expand Down
3 changes: 1 addition & 2 deletions cosmos/dbt/parser/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,13 @@

import logging
import re
from typing import List, Tuple, TYPE_CHECKING
from typing import TYPE_CHECKING, List, Tuple

if TYPE_CHECKING:
from dbt.cli.main import dbtRunnerResult

from cosmos.hooks.subprocess import FullOutputSubprocessResult


DBT_NO_TESTS_MSG = "Nothing to do"
DBT_WARN_MSG = "WARN"

Expand Down
2 changes: 1 addition & 1 deletion cosmos/dbt/parser/project.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@

from __future__ import annotations

import os
import ast
import os
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
Expand Down
7 changes: 4 additions & 3 deletions cosmos/dbt/project.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
from __future__ import annotations

from pathlib import Path
import shutil
import os
from cosmos.constants import DBT_LOG_DIR_NAME, DBT_TARGET_DIR_NAME, DBT_PARTIAL_PARSE_FILE_NAME
import shutil
from contextlib import contextmanager
from pathlib import Path
from typing import Generator

from cosmos.constants import DBT_LOG_DIR_NAME, DBT_PARTIAL_PARSE_FILE_NAME, DBT_TARGET_DIR_NAME


def create_symlinks(project_path: Path, tmp_dir: Path, ignore_dbt_packages: bool) -> None:
"""Helper function to create symlinks to the dbt project files."""
Expand Down
1 change: 1 addition & 0 deletions cosmos/dbt/selector.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from __future__ import annotations

import copy
import re
from collections import defaultdict
Expand Down
2 changes: 1 addition & 1 deletion cosmos/hooks/subprocess.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@
import contextlib
import os
import signal
from typing import NamedTuple
from subprocess import PIPE, STDOUT, Popen
from tempfile import TemporaryDirectory, gettempdir
from typing import NamedTuple

from airflow.hooks.base import BaseHook

Expand Down
2 changes: 1 addition & 1 deletion cosmos/log.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
from __future__ import annotations

import logging

from airflow.configuration import conf
from airflow.utils.log.colored_log import CustomTTYColoredFormatter


LOG_FORMAT: str = (
"[%(blue)s%(asctime)s%(reset)s] "
"{%(blue)s%(filename)s:%(reset)s%(lineno)d} "
Expand Down
2 changes: 1 addition & 1 deletion cosmos/operators/__init__.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
from .local import DbtBuildLocalOperator as DbtBuildOperator
from .local import DbtDepsLocalOperator as DbtDepsOperator
from .local import DbtDocsAzureStorageLocalOperator as DbtDocsAzureStorageOperator
from .local import DbtDocsGCSLocalOperator as DbtDocsGCSOperator
from .local import DbtDocsLocalOperator as DbtDocsOperator
from .local import DbtDocsS3LocalOperator as DbtDocsS3Operator
from .local import DbtDocsGCSLocalOperator as DbtDocsGCSOperator
from .local import DbtLSLocalOperator as DbtLSOperator
from .local import DbtRunLocalOperator as DbtRunOperator
from .local import DbtRunOperationLocalOperator as DbtRunOperationOperator
Expand Down
6 changes: 3 additions & 3 deletions cosmos/operators/azure_container_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,17 @@
from typing import Any, Callable, Sequence

from airflow.utils.context import Context
from cosmos.config import ProfileConfig

from cosmos.config import ProfileConfig
from cosmos.log import get_logger
from cosmos.operators.base import (
AbstractDbtBaseOperator,
DbtLSMixin,
DbtRunMixin,
DbtRunOperationMixin,
DbtSeedMixin,
DbtSnapshotMixin,
DbtTestMixin,
DbtLSMixin,
DbtRunOperationMixin,
)

logger = get_logger(__name__)
Expand Down
3 changes: 1 addition & 2 deletions cosmos/operators/base.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
from __future__ import annotations

import os
from typing import Any, Sequence, Tuple
from abc import ABCMeta, abstractmethod
from typing import Any, Sequence, Tuple

import yaml
from airflow.models.baseoperator import BaseOperator
Expand All @@ -12,7 +12,6 @@
from cosmos.dbt.executable import get_system_dbt
from cosmos.log import get_logger


logger = get_logger(__name__)


Expand Down
4 changes: 2 additions & 2 deletions cosmos/operators/docker.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,12 +8,12 @@
from cosmos.operators.base import (
AbstractDbtBaseOperator,
DbtBuildMixin,
DbtLSMixin,
DbtRunMixin,
DbtRunOperationMixin,
DbtSeedMixin,
DbtSnapshotMixin,
DbtTestMixin,
DbtLSMixin,
DbtRunOperationMixin,
)

logger = get_logger(__name__)
Expand Down
11 changes: 5 additions & 6 deletions cosmos/operators/kubernetes.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,24 +3,23 @@
from os import PathLike
from typing import Any, Callable, Sequence

from airflow.models import TaskInstance
from airflow.utils.context import Context, context_merge

from cosmos.log import get_logger
from cosmos.config import ProfileConfig
from cosmos.dbt.parser.output import extract_log_issues
from cosmos.log import get_logger
from cosmos.operators.base import (
AbstractDbtBaseOperator,
DbtBuildMixin,
DbtLSMixin,
DbtRunMixin,
DbtRunOperationMixin,
DbtSeedMixin,
DbtSnapshotMixin,
DbtTestMixin,
DbtLSMixin,
DbtRunOperationMixin,
)

from airflow.models import TaskInstance
from cosmos.dbt.parser.output import extract_log_issues

DBT_NO_TESTS_MSG = "Nothing to do"
DBT_WARN_MSG = "WARN"

Expand Down
Loading

0 comments on commit 8b3441e

Please sign in to comment.