feat(ingest): remove dep on termcolor

Centralize on click.style and click.echo. This ensures colors don't get garbled on non-color terminals, and also ensures that we consistently print to stderr where we're supposed to.
datahub-project · Oct 29, 2024 · a5e89d2 · a5e89d2
1 parent ea7bdca
commit a5e89d2
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 45 deletions.
diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
@@ -42,7 +42,6 @@
     "python-dateutil>=2.8.0",
     "tabulate",
     "progressbar2",
-    "termcolor>=1.0.0",
     "psutil>=5.8.0",
     "Deprecated",
     "humanfriendly",
@@ -546,7 +545,6 @@
     "types-pyOpenSSL",
     "types-click-spinner>=0.1.13.1",
     "types-ujson>=5.2.0",
-    "types-termcolor>=1.0.0",
     "types-Deprecated",
     "types-protobuf>=4.21.0.1",
     "sqlalchemy2-stubs",

diff --git a/metadata-ingestion/src/datahub/cli/specific/dataproduct_cli.py b/metadata-ingestion/src/datahub/cli/specific/dataproduct_cli.py
@@ -363,7 +363,7 @@ def remove_owner(urn: str, owner_urn: str) -> None:
     with get_default_graph() as graph:
         _abort_if_non_existent_urn(graph, urn, "remove owners")
         for mcp in dataproduct_patcher.build():
-            print(json.dumps(mcp.to_obj()))
+            click.echo(json.dumps(mcp.to_obj()))
             graph.emit(mcp)
 
 

diff --git a/metadata-ingestion/src/datahub/cli/timeline_cli.py b/metadata-ingestion/src/datahub/cli/timeline_cli.py
@@ -6,7 +6,6 @@
 
 import click
 from requests import Response
-from termcolor import colored
 
 from datahub.emitter.mce_builder import dataset_urn_to_key, schema_field_urn_to_key
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
@@ -44,14 +43,14 @@ def pretty_id(id: Optional[str]) -> str:
             assert schema_field_key is not None
             field_path = schema_field_key.fieldPath
 
-            return f"{colored('field','cyan')}:{colored(pretty_field_path(field_path),'white')}"
+            return f"{click.style('field', fg='cyan')}:{click.style(pretty_field_path(field_path), fg='white')}"
     if id.startswith("[version=2.0]"):
-        return f"{colored('field','cyan')}:{colored(pretty_field_path(id),'white')}"
+        return f"{click.style('field', fg='cyan')}:{click.style(pretty_field_path(id), fg='white')}"
 
     if id.startswith("urn:li:dataset"):
         dataset_key = dataset_urn_to_key(id)
         if dataset_key:
-            return f"{colored('dataset','cyan')}:{colored(dataset_key.platform[len('urn:li:dataPlatform:'):],'white')}:{colored(dataset_key.name,'white')}"
+            return f"{click.style('dataset', fg='cyan')}:{click.style(dataset_key.platform[len('urn:li:dataPlatform:'):], fg='white')}:{click.style(dataset_key.name, fg='white')}"
     # failed to prettify, return original
     return id
 
@@ -196,8 +195,8 @@ def timeline(
                 else "red"
             )
 
-            print(
-                f"{colored(change_instant,'cyan')} - {colored(change_txn['semVer'],change_color)}"
+            click.echo(
+                f"{click.style(change_instant, fg='cyan')} - {click.style(change_txn['semVer'], fg=change_color)}"
             )
             if change_txn["changeEvents"] is not None:
                 for change_event in change_txn["changeEvents"]:
@@ -216,8 +215,8 @@ def timeline(
                         or change_event.get("entityUrn")
                         or ""
                     )
-                    print(
-                        f"\t{colored(change_event.get('changeType') or change_event.get('operation'),event_change_color)} {change_event.get('category')} {target_string} {element_string}: {change_event['description']}"
+                    click.echo(
+                        f"\t{click.style(change_event.get('changeType') or change_event.get('operation'), fg=event_change_color)} {change_event.get('category')} {target_string} {element_string}: {change_event['description']}"
                     )
     else:
         click.echo(

diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
@@ -647,11 +647,11 @@ def _get_text_color(self, running: bool, failures: bool, warnings: bool) -> str:
             return "cyan"
         else:
             if failures:
-                return "bright_red"
+                return "red"
             elif warnings:
-                return "bright_yellow"
+                return "yellow"
             else:
-                return "bright_green"
+                return "green"
 
     def has_failures(self) -> bool:
         return bool(
@@ -674,7 +674,7 @@ def pretty_print_summary(
         else:
             click.echo()
             click.secho("Cli report:", bold=True)
-            click.secho(self.cli_report.as_string())
+            click.echo(self.cli_report.as_string())
             click.secho(f"Source ({self.source_type}) report:", bold=True)
             click.echo(self.source.get_report().as_string())
             click.secho(f"Sink ({self.sink_type}) report:", bold=True)

diff --git a/metadata-ingestion/src/datahub/upgrade/upgrade.py b/metadata-ingestion/src/datahub/upgrade/upgrade.py
@@ -1,15 +1,14 @@
 import asyncio
 import contextlib
 import logging
-import sys
 from datetime import datetime, timedelta, timezone
 from functools import wraps
 from typing import Any, Callable, Optional, Tuple, TypeVar
 
+import click
 import humanfriendly
 from packaging.version import Version
 from pydantic import BaseModel
-from termcolor import colored
 
 from datahub import __version__
 from datahub.cli.config_utils import load_client_config
@@ -277,8 +276,8 @@ def maybe_print_upgrade_message(  # noqa: C901
         if not version_stats:
             log.debug("No version stats found")
             return
-        else:
-            log.debug(f"Version stats found: {version_stats}")
+
+        log.debug(f"Version stats found: {version_stats}")
         current_release_date = version_stats.client.current.release_date
         latest_release_date = (
             version_stats.client.latest.release_date
@@ -325,50 +324,54 @@ def maybe_print_upgrade_message(  # noqa: C901
     if client_server_compat < 0:
         with contextlib.suppress(Exception):
             assert version_stats
-            print(
-                colored("❗Client-Server Incompatible❗", "yellow"),
-                colored(
+            click.echo(
+                click.style("❗Client-Server Incompatible❗", fg="yellow")
+                + " "
+                + click.style(
                     f"Your client version {version_stats.client.current.version} is newer than your server version {version_stats.server.current.version}. Downgrading the cli to {version_stats.server.current.version} is recommended.\n",
-                    "cyan",
-                ),
-                colored(
+                    fg="cyan",
+                )
+                + click.style(
                     f"➡️ Downgrade via `\"pip install 'acryl-datahub=={version_stats.server.current.version}'\"",
-                    "cyan",
-                ),
+                    fg="cyan",
+                )
             )
     elif client_server_compat > 0:
         with contextlib.suppress(Exception):
             assert version_stats
-            print(
-                colored("❗Client-Server Incompatible❗", "red"),
-                colored(
+            click.echo(
+                click.style("❗Client-Server Incompatible❗", fg="red")
+                + " "
+                + click.style(
                     f"Your client version {version_stats.client.current.version} is older than your server version {version_stats.server.current.version}. Upgrading the cli to {version_stats.server.current.version} is recommended.\n",
-                    "cyan",
-                ),
-                colored(
+                    fg="cyan",
+                )
+                + click.style(
                     f"➡️  Upgrade via \"pip install 'acryl-datahub=={version_stats.server.current.version}'\"",
-                    "cyan",
-                ),
+                    fg="cyan",
+                )
             )
     elif client_server_compat == 0 and encourage_cli_upgrade:
         with contextlib.suppress(Exception):
-            print(
-                colored("💡 Upgrade cli!", "yellow"),
-                colored(
+            click.echo(
+                click.style("💡 Upgrade cli!", fg="yellow")
+                + " "
+                + click.style(
                     f"You seem to be running an old version of datahub cli: {current_version} {get_days(current_release_date)}. Latest version is {latest_version} {get_days(latest_release_date)}.\nUpgrade via \"pip install -U 'acryl-datahub'\"",
-                    "cyan",
-                ),
+                    fg="cyan",
+                )
             )
     elif encourage_quickstart_upgrade:
         try:
             assert version_stats
-            print(
-                colored("💡 Upgrade available!", "yellow"),
-                colored(
+            click.echo(
+                click.style("💡 Upgrade available!", fg="yellow")
+                + " "
+                + click.style(
                     f'You seem to be running a slightly old quickstart image {get_days(version_stats.server.current.release_date)}. Run "datahub docker quickstart" to get the latest updates without losing any data!',
-                    "cyan",
+                    fg="cyan",
                 ),
-                file=sys.stderr,
+                err=True,
             )
         except Exception as e:
             log.debug(f"Failed to suggest quickstart upgrade due to {e}")