From 35eee7821a7456b2ea0347706c1827f9b027f4ef Mon Sep 17 00:00:00 2001 From: "Wei-Chun, Chang" Date: Tue, 7 Nov 2023 11:25:56 +0800 Subject: [PATCH] Wrap ruamel yaml library import due to type support Signed-off-by: Wei-Chun, Chang --- piperider_cli/__init__.py | 5 +- piperider_cli/configuration.py | 10 ++-- piperider_cli/dbtutil.py | 5 +- piperider_cli/event/__init__.py | 5 +- piperider_cli/initializer.py | 5 +- piperider_cli/recipes/__init__.py | 5 +- piperider_cli/yaml/__init__.py | 60 ------------------- piperider_cli/yaml/yaml_less_type/__init__.py | 54 +++++++++++++++++ piperider_cli/yaml/yaml_rich_type/__init__.py | 60 +++++++++++++++++++ 9 files changed, 140 insertions(+), 69 deletions(-) create mode 100644 piperider_cli/yaml/yaml_less_type/__init__.py create mode 100644 piperider_cli/yaml/yaml_rich_type/__init__.py diff --git a/piperider_cli/__init__.py b/piperider_cli/__init__.py index 4be7012ff..0d18b816e 100644 --- a/piperider_cli/__init__.py +++ b/piperider_cli/__init__.py @@ -9,7 +9,10 @@ from dateutil import tz from rich.console import Console -from piperider_cli import yaml as pyml +try: + from piperider_cli.yaml import yaml_rich_type as pyml +except ImportError: + from piperider_cli.yaml import yaml_less_type as pyml PIPERIDER_USER_HOME = os.path.expanduser('~/.piperider') if os.access(os.path.expanduser('~/'), os.W_OK) is False: diff --git a/piperider_cli/configuration.py b/piperider_cli/configuration.py index 714305f1a..215907554 100644 --- a/piperider_cli/configuration.py +++ b/piperider_cli/configuration.py @@ -13,7 +13,6 @@ from rich.console import Console from piperider_cli import raise_exception_when_directory_not_writable -from piperider_cli.yaml import round_trip_load_yaml, safe_load_yaml from piperider_cli.cli_utils import DbtUtil from piperider_cli.datasource import DATASOURCE_PROVIDERS, DataSource from piperider_cli.datasource.unsupported import UnsupportedDataSource @@ -23,7 +22,10 @@ PipeRiderInvalidDataSourceError, \ DbtProjectNotFoundError, \ DbtProfileNotFoundError -from piperider_cli import yaml as pyml +try: + from piperider_cli.yaml import yaml_rich_type as pyml +except ImportError: + from piperider_cli.yaml import yaml_less_type as pyml # ref: https://docs.getdbt.com/dbt-cli/configure-your-profile DBT_PROFILES_DIR_DEFAULT = '~/.dbt/' @@ -488,7 +490,7 @@ def _load(cls, piperider_config_path=None, dbt_profile: str = None, dbt_target: credentials = None piperider_config_path = piperider_config_path or FileSystem.PIPERIDER_CONFIG_PATH - config = safe_load_yaml(piperider_config_path) + config = pyml.safe_load_yaml(piperider_config_path) if config is None: raise PipeRiderConfigError(piperider_config_path) @@ -602,7 +604,7 @@ def flush_datasource(self, path): :return: """ - config = round_trip_load_yaml(path) + config = pyml.round_trip_load_yaml(path) if config is None: raise PipeRiderConfigError(path) diff --git a/piperider_cli/dbtutil.py b/piperider_cli/dbtutil.py index 1f3ac88c3..caf28e942 100644 --- a/piperider_cli/dbtutil.py +++ b/piperider_cli/dbtutil.py @@ -10,7 +10,10 @@ from jinja2 import UndefinedError from rich.console import Console from rich.table import Table -from piperider_cli import yaml as pyml +try: + from piperider_cli.yaml import yaml_rich_type as pyml +except ImportError: + from piperider_cli.yaml import yaml_less_type as pyml from piperider_cli import load_jinja_template, load_jinja_string_template from piperider_cli.dbt.list_task import load_manifest, list_resources_unique_id_from_manifest, load_full_manifest diff --git a/piperider_cli/event/__init__.py b/piperider_cli/event/__init__.py index 1cc55d38d..77d7a83db 100644 --- a/piperider_cli/event/__init__.py +++ b/piperider_cli/event/__init__.py @@ -5,7 +5,10 @@ import sentry_sdk from rich.console import Console -from piperider_cli import yaml as pyml +try: + from piperider_cli.yaml import yaml_rich_type as pyml +except ImportError: + from piperider_cli.yaml import yaml_less_type as pyml from piperider_cli import PIPERIDER_USER_HOME, PIPERIDER_USER_PROFILE, is_executed_manually from piperider_cli.event.collector import Collector diff --git a/piperider_cli/initializer.py b/piperider_cli/initializer.py index 3d92ed5c8..594703759 100644 --- a/piperider_cli/initializer.py +++ b/piperider_cli/initializer.py @@ -8,7 +8,10 @@ from rich.table import Table from piperider_cli import clone_directory -from piperider_cli.yaml import safe_load_yaml +try: + from piperider_cli.yaml.yaml_rich_type import safe_load_yaml +except ImportError: + from piperider_cli.yaml.yaml_less_type import safe_load_yaml from piperider_cli.configuration import Configuration, FileSystem from piperider_cli.datasource import DataSource, FANCY_USER_INPUT diff --git a/piperider_cli/recipes/__init__.py b/piperider_cli/recipes/__init__.py index 5c749c44d..b533c5a6e 100644 --- a/piperider_cli/recipes/__init__.py +++ b/piperider_cli/recipes/__init__.py @@ -11,7 +11,10 @@ import piperider_cli.dbtutil as dbtutil from piperider_cli import get_run_json_path, load_jinja_template, load_json -from piperider_cli import yaml as pyml +try: + from piperider_cli.yaml import yaml_rich_type as pyml +except ImportError: + from piperider_cli.yaml import yaml_less_type as pyml from piperider_cli.configuration import Configuration, FileSystem from piperider_cli.dbt import dbt_version from piperider_cli.error import RecipeConfigException diff --git a/piperider_cli/yaml/__init__.py b/piperider_cli/yaml/__init__.py index ccc15d2bd..e69de29bb 100644 --- a/piperider_cli/yaml/__init__.py +++ b/piperider_cli/yaml/__init__.py @@ -1,60 +0,0 @@ -from pathlib import Path -from typing import Any, Callable, Optional, Union - -from ruamel import yaml -from ruamel.yaml import StreamTextType, StreamType, VersionType -from ruamel.yaml import CommentedMap as _cm, CommentedSeq as _cs - -_yaml = yaml.YAML() - -CommentedMap = _cm -CommentedSeq = _cs -YAMLError = yaml.YAMLError - - -def load(stream: Union[Path, StreamTextType]) -> Any: - return _yaml.load(stream) - - -def allow_duplicate_keys_loader() -> Callable: - yml = yaml.YAML() - yml.allow_duplicate_keys = True - return yml.load - - -def safe_load(stream: StreamTextType, version: Optional[VersionType] = None) -> Any: - return yaml.safe_load(stream, version) - - -def dump( - data: Union[Path, StreamType], stream: Any = None, *, transform: Any = None -) -> Any: - return _yaml.dump(data, stream, transform=transform) - - -def safe_load_yaml(file_path): - try: - with open(file_path, 'r') as f: - payload = safe_load(f) - except yaml.YAMLError as e: - print(e) - return None - except FileNotFoundError: - return None - return payload - - -def round_trip_load_yaml(file_path): - with open(file_path, 'r') as f: - try: - payload = load(f) - except yaml.YAMLError as e: - print(e) - return None - return payload - - -def round_trip_dump( - data: Any, - stream: Optional[StreamType] = None): - return yaml.round_trip_dump(data, stream) diff --git a/piperider_cli/yaml/yaml_less_type/__init__.py b/piperider_cli/yaml/yaml_less_type/__init__.py new file mode 100644 index 000000000..b0c6e3806 --- /dev/null +++ b/piperider_cli/yaml/yaml_less_type/__init__.py @@ -0,0 +1,54 @@ +from typing import Any, Callable + +from ruamel import yaml +from ruamel.yaml import CommentedMap as _cm, CommentedSeq as _cs + +_yaml = yaml.YAML() + +CommentedMap = _cm +CommentedSeq = _cs +YAMLError = yaml.YAMLError + + +def load(stream) -> Any: + return _yaml.load(stream) + + +def allow_duplicate_keys_loader() -> Callable: + yml = yaml.YAML() + yml.allow_duplicate_keys = True + return yml.load + + +def safe_load(stream, version=None) -> Any: + return yaml.safe_load(stream, version) + + +def dump(data, stream=None, *, transform=None) -> Any: + return _yaml.dump(data, stream, transform=transform) + + +def safe_load_yaml(file_path): + try: + with open(file_path, 'r') as f: + payload = safe_load(f) + except yaml.YAMLError as e: + print(e) + return None + except FileNotFoundError: + return None + return payload + + +def round_trip_load_yaml(file_path): + with open(file_path, 'r') as f: + try: + payload = load(f) + except yaml.YAMLError as e: + print(e) + return None + return payload + + +def round_trip_dump(data, stream=None): + return yaml.round_trip_dump(data, stream) diff --git a/piperider_cli/yaml/yaml_rich_type/__init__.py b/piperider_cli/yaml/yaml_rich_type/__init__.py new file mode 100644 index 000000000..ccc15d2bd --- /dev/null +++ b/piperider_cli/yaml/yaml_rich_type/__init__.py @@ -0,0 +1,60 @@ +from pathlib import Path +from typing import Any, Callable, Optional, Union + +from ruamel import yaml +from ruamel.yaml import StreamTextType, StreamType, VersionType +from ruamel.yaml import CommentedMap as _cm, CommentedSeq as _cs + +_yaml = yaml.YAML() + +CommentedMap = _cm +CommentedSeq = _cs +YAMLError = yaml.YAMLError + + +def load(stream: Union[Path, StreamTextType]) -> Any: + return _yaml.load(stream) + + +def allow_duplicate_keys_loader() -> Callable: + yml = yaml.YAML() + yml.allow_duplicate_keys = True + return yml.load + + +def safe_load(stream: StreamTextType, version: Optional[VersionType] = None) -> Any: + return yaml.safe_load(stream, version) + + +def dump( + data: Union[Path, StreamType], stream: Any = None, *, transform: Any = None +) -> Any: + return _yaml.dump(data, stream, transform=transform) + + +def safe_load_yaml(file_path): + try: + with open(file_path, 'r') as f: + payload = safe_load(f) + except yaml.YAMLError as e: + print(e) + return None + except FileNotFoundError: + return None + return payload + + +def round_trip_load_yaml(file_path): + with open(file_path, 'r') as f: + try: + payload = load(f) + except yaml.YAMLError as e: + print(e) + return None + return payload + + +def round_trip_dump( + data: Any, + stream: Optional[StreamType] = None): + return yaml.round_trip_dump(data, stream)