From 93920baaec2f1351600b7deb2fecbc015b9cc9b0 Mon Sep 17 00:00:00 2001
From: "Mr. Outis" <mroutis@protonmail.com>
Date: Mon, 25 Feb 2019 11:31:43 -0600
Subject: [PATCH] logger: Format progress bar interruption

Fix #1602
---
 dvc/logger.py          |  4 ++++
 dvc/progress.py        | 17 ++++++++++++++++-
 dvc/prompt.py          |  3 +++
 tests/unit/progress.py | 20 ++++++++++++++++++++
 4 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/progress.py

diff --git a/dvc/logger.py b/dvc/logger.py
index abc02aed4..69e16e1dc 100644
--- a/dvc/logger.py
+++ b/dvc/logger.py
@@ -4,6 +4,7 @@ from __future__ import unicode_literals
 
 from dvc.exceptions import DvcException
 from dvc.utils.compat import str
+from dvc.progress import progress_aware
 
 import re
 import sys
@@ -15,6 +16,7 @@ from contextlib import contextmanager
 import colorama
 
 
+@progress_aware
 def info(message):
     """Prints an info message."""
     logger.info(message)
@@ -29,6 +31,7 @@ def debug(message):
     logger.debug(out)
 
 
+@progress_aware
 def warning(message, parse_exception=False):
     """Prints a warning message."""
     prefix = colorize("Warning", color="yellow")
@@ -47,6 +50,7 @@ def warning(message, parse_exception=False):
     logger.warning(out)
 
 
+@progress_aware
 def error(message=None):
     """Prints an error message."""
     prefix = colorize("Error", color="red")
diff --git a/dvc/progress.py b/dvc/progress.py
index 7b6544238..05315f8f7 100644
--- a/dvc/progress.py
+++ b/dvc/progress.py
@@ -7,7 +7,6 @@ from dvc.utils.compat import str
 
 import sys
 import threading
-import dvc.logger as logger
 
 
 class Progress(object):
@@ -100,6 +99,8 @@ class Progress(object):
 
     @staticmethod
     def _print(*args, **kwargs):
+        import dvc.logger as logger
+
         if logger.is_quiet():
             return
 
@@ -116,4 +117,18 @@ class Progress(object):
         self._lock.release()
 
 
+def progress_aware(f):
+    """ Decorator to add a new line if progress bar hasn't finished  """
+    from functools import wraps
+
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        if not progress.is_finished:
+            progress._print()
+
+        return f(*args, **kwargs)
+
+    return wrapper
+
+
 progress = Progress()  # pylint: disable=invalid-name
diff --git a/dvc/prompt.py b/dvc/prompt.py
index 898eb921b..b4fa15c31 100644
--- a/dvc/prompt.py
+++ b/dvc/prompt.py
@@ -6,6 +6,7 @@ from __future__ import print_function
 import sys
 from getpass import getpass
 
+from dvc.progress import progress_aware
 from dvc.utils.compat import input
 
 
@@ -32,6 +33,7 @@ def _ask(prompt, limited_to=None):
         )
 
 
+@progress_aware
 def confirm(statement):
     """Ask the user for confirmation about the specified statement.
 
@@ -46,6 +48,7 @@ def confirm(statement):
     return answer and answer.startswith("y")
 
 
+@progress_aware
 def password(statement):
     """Ask the user for a password.
 
diff --git a/tests/unit/progress.py b/tests/unit/progress.py
new file mode 100644
index 000000000..3a2974147
--- /dev/null
+++ b/tests/unit/progress.py
@@ -0,0 +1,20 @@
+import mock
+from unittest import TestCase
+
+from dvc.progress import progress, progress_aware
+
+
+class TestProgressAware(TestCase):
+    @mock.patch("sys.stdout.isatty", return_value=True)
+    @mock.patch("dvc.progress.Progress._print")
+    def test(self, mock_print, mock_isatty):
+        function = progress_aware(lambda: None)
+
+        function()
+        mock_print.assert_not_called()
+
+        progress.update_target("testing", 0, 100)
+        function()
+        mock_print.assert_called()
+
+        progress.finish_target("testing")