Wrap ruamel yaml library import due to type support

Signed-off-by: Wei-Chun, Chang <[email protected]>

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:

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)
 

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

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

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

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

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)

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)

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)