From 8312eb9b424ff88794e32220f6e5c1be84532122 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Wed, 15 Mar 2023 17:42:45 +0545
Subject: [PATCH] ui: disable colorama on write_json when highlighting is
 disabled (#9186)

On a large json output (say ~300MB), this saves about ~6 seconds of time.
---
 dvc/ui/__init__.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py
index e11abeb8a..5e81a2e75 100644
--- a/dvc/ui/__init__.py
+++ b/dvc/ui/__init__.py
@@ -13,6 +13,8 @@ from typing import (
     Union,
 )
 
+import colorama
+
 from dvc.utils.objects import cached_property
 
 if TYPE_CHECKING:
@@ -27,6 +29,15 @@ if TYPE_CHECKING:
     from dvc.ui.table import Headers, Styles, TableData
 
 
+@contextmanager
+def disable_colorama():
+    colorama.deinit()
+    try:
+        yield
+    finally:
+        colorama.reinit()
+
+
 class Formatter:
     def __init__(
         self, theme: Optional[Dict] = None, defaults: Optional[Dict] = None
@@ -121,7 +132,10 @@ class Console:
             sort_keys=sort_keys,
         )
         if not highlight:
-            return self.write(json.text, stderr=stderr)
+            # we don't need colorama to try to strip ansi codes
+            # when highlighting is disabled
+            with disable_colorama():
+                return self.write(json.text, stderr=stderr)
         return self.rich_print(json, stderr=stderr, soft_wrap=True)
 
     def rich_print(