Feature/progress bar print (#71)

Author: Juliya Smith

CHANGELOG.md

@@ -51,6 +51,8 @@ how a consumer would use the library (e.g. adding unit tests, updating documenta
 - Warning message printed when ctrl-c is encountered in the middle of an operation that could cause incorrect checkpoint
     state, a second ctrl-c is required to quit while that operation is ongoing. 
 
+- A progress bar that displays during bulk commands.
+
 ### Fixed
 
 - Fixed bug in bulk commands where value-less fields in csv files were treated as empty strings instead of None.

src/code42cli/bulk.py

@@ -1,11 +1,13 @@
-import os
-import inspect
-import csv
+import os, inspect
 
 from code42cli.compat import open, str
 from code42cli.worker import Worker
 from code42cli.logger import get_main_cli_logger
 from code42cli.args import SDK_ARG_NAME, PROFILE_ARG_NAME
+from code42cli.progress_bar import ProgressBar
+
+
+_logger = get_main_cli_logger()
 
 
 class BulkCommandType(object):
@@ -29,7 +31,7 @@ def generate_template(handler, path=None):
     ]
 
     if len(args) <= 1:
-        get_main_cli_logger().print_info(
+        _logger.print_info(
             u"A blank file was generated because there are no csv headers needed for this command. "
             u"Simply enter one {} per line.".format(args[0])
         )
@@ -45,31 +47,28 @@ def _write_template_file(path, columns=None):
             new_file.write(u",".join(columns))
 
 
-def run_bulk_process(file_path, row_handler, reader=None):
+def run_bulk_process(row_handler, reader):
     """Runs a bulk process.
     
     Args: 
-        file_path (str or unicode): The path to the file feeding the data for the bulk process.
         row_handler (callable): A callable that you define to process values from the row as 
             either *args or **kwargs.
         reader: (CSVReader or FlatFileReader, optional): A generator that reads rows and yields data into 
             `row_handler`. If None, it will use a CSVReader. Defaults to None.
     """
-    reader = reader or CSVReader()
-    processor = _create_bulk_processor(file_path, row_handler, reader)
+    processor = _create_bulk_processor(row_handler, reader)
     processor.run()
 
 
-def _create_bulk_processor(file_path, row_handler, reader):
+def _create_bulk_processor(row_handler, reader):
     """A factory method to create the bulk processor, useful for testing purposes."""
-    return BulkProcessor(file_path, row_handler, reader)
+    return BulkProcessor(row_handler, reader)
 
 
 class BulkProcessor(object):
     """A class for bulk processing a file. 
     
     Args:
-        file_path (str or unicode): The path to the file for processing.
         row_handler (callable): A callable that you define to process values from the row as 
             either *args or **kwargs. For example, if it's a csv file with header `prop_a,prop_b` 
             and first row `1,test`, then `row_handler` should receive kwargs 
@@ -78,19 +77,22 @@ class BulkProcessor(object):
         reader (CSVReader or FlatFileReader): A generator that reads rows and yields data into `row_handler`.
     """
 
-    def __init__(self, file_path, row_handler, reader):
-        self.file_path = file_path
+    def __init__(self, row_handler, reader, worker=None, progress_bar=None):
+        total = reader.get_rows_count()
+        self.file_path = reader.file_path
         self._row_handler = row_handler
         self._reader = reader
-        self.__worker = Worker(5)
+        self.__worker = worker or Worker(5, total)
+        self._stats = self.__worker.stats
+        self._progress_bar = progress_bar or ProgressBar(total)
 
     def run(self):
         """Processes the csv file specified in the ctor, calling `self.row_handler` on each row."""
         with open(self.file_path, newline=u"", encoding=u"utf8") as bulk_file:
             for row in self._reader(bulk_file=bulk_file):
                 self._process_row(row)
             self.__worker.wait()
-        self._print_result()
+        self._print_results()
 
     def _process_row(self, row):
         if isinstance(row, dict):
@@ -104,35 +106,20 @@ def _process_csv_row(self, row):
         row.pop(None, None)
         row_values = {key: val if val != u"" else None for key, val in row.items()}
         self.__worker.do_async(
-            lambda *args, **kwargs: self._row_handler(*args, **kwargs), **row_values
+            lambda *args, **kwargs: self._handle_row(*args, **kwargs), **row_values
         )
 
     def _process_flat_file_row(self, row):
         if row:
-            self.__worker.do_async(lambda *args, **kwargs: self._row_handler(*args, **kwargs), row)
-
-    def _print_result(self):
-        stats = self.__worker.stats
-        successes = stats.total - stats.total_errors
-        logger = get_main_cli_logger()
-        logger.print_and_log_info(
-            u"{} processed successfully out of {}.".format(successes, stats.total)
-        )
-        if stats.total_errors:
-            logger.print_errors_occurred_message()
-
+            self.__worker.do_async(lambda *args, **kwargs: self._handle_row(*args, **kwargs), row)
 
-class CSVReader(object):
-    """A generator that yields header keys mapped to row values from a csv file."""
+    def _handle_row(self, *args, **kwargs):
+        message = str(self._stats)
+        self._progress_bar.update(self._stats.total_processed, message)
+        self._row_handler(*args, **kwargs)
 
-    def __call__(self, *args, **kwargs):
-        for row in csv.DictReader(kwargs.get(u"bulk_file")):
-            yield row
-
-
-class FlatFileReader(object):
-    """A generator that yields a single-value per row from a file."""
-
-    def __call__(self, *args, **kwargs):
-        for row in kwargs[u"bulk_file"]:
-            yield row
+    def _print_results(self):
+        self._progress_bar.clear_bar_and_print_final(str(self._stats))
+        if self._stats.total_errors:
+            logger = get_main_cli_logger()
+            logger.print_errors_occurred_message()

src/code42cli/cmds/alerts/rules/user_rule.py

@@ -4,7 +4,8 @@
 
 from code42cli.errors import InvalidRuleTypeError
 from code42cli.util import format_to_table, find_format_width
-from code42cli.bulk import run_bulk_process, CSVReader
+from code42cli.bulk import run_bulk_process
+from code42cli.file_readers import create_csv_reader
 from code42cli.logger import get_main_cli_logger
 from code42cli.cmds.detectionlists import get_user_id
 from code42cli.cmds.alerts.rules.enums import AlertRuleTypes
@@ -74,17 +75,13 @@ def get_rules(sdk, profile):
 
 
 def add_bulk_users(sdk, profile, file_name):
-    run_bulk_process(
-        file_name, lambda rule_id, username: add_user(sdk, profile, rule_id, username), CSVReader()
-    )
+    reader = create_csv_reader(file_name)
+    run_bulk_process(lambda rule_id, username: add_user(sdk, profile, rule_id, username), reader)
 
 
 def remove_bulk_users(sdk, profile, file_name):
-    run_bulk_process(
-        file_name,
-        lambda rule_id, username: remove_user(sdk, profile, rule_id, username),
-        CSVReader(),
-    )
+    reader = create_csv_reader(file_name)
+    run_bulk_process(lambda rule_id, username: remove_user(sdk, profile, rule_id, username), reader)
 
 
 def show_rule(sdk, profile, rule_id):

src/code42cli/cmds/detectionlists/__init__.py

@@ -1,10 +1,9 @@
 from py42.exceptions import Py42BadRequestError
 
-from code42cli.compat import str
+from code42cli.bulk import generate_template, run_bulk_process, BulkCommandType
+from code42cli.file_readers import create_csv_reader, create_flat_file_reader
+from code42cli.errors import UserAlreadyAddedError, UserDoesNotExistError, UnknownRiskTagError
 from code42cli.cmds.detectionlists.commands import DetectionListCommandFactory
-from code42cli.bulk import generate_template, run_bulk_process, CSVReader, FlatFileReader
-from code42cli.errors import UserAlreadyAddedError, UnknownRiskTagError, UserDoesNotExistError
-from code42cli.bulk import BulkCommandType
 from code42cli.cmds.detectionlists.enums import DetectionLists, DetectionListUserKeys, RiskTags
 
 
@@ -123,9 +122,8 @@ def bulk_add_employees(self, sdk, profile, csv_file):
             profile (Code42Profile): The profile under which to execute this command.
             csv_file (str or unicode): The path to the csv file containing rows of users.
         """
-        run_bulk_process(
-            csv_file, lambda **kwargs: self._add_employee(sdk, profile, **kwargs), CSVReader()
-        )
+        reader = create_csv_reader(csv_file)
+        run_bulk_process(lambda **kwargs: self._add_employee(sdk, profile, **kwargs), reader)
 
     def bulk_remove_employees(self, sdk, profile, users_file):
         """Takes a flat file with each row containing a username and removes them all from the 
@@ -136,10 +134,9 @@ def bulk_remove_employees(self, sdk, profile, users_file):
             profile (Code42Profile): The profile under which to execute this command.
             users_file (str or unicode): The path to the file containing rows of user names.
         """
+        reader = create_flat_file_reader(users_file)
         run_bulk_process(
-            users_file,
-            lambda *args, **kwargs: self._remove_employee(sdk, profile, *args, **kwargs),
-            FlatFileReader(),
+            lambda *args, **kwargs: self._remove_employee(sdk, profile, *args, **kwargs), reader
         )
 
     def _add_employee(self, sdk, profile, **kwargs):

src/code42cli/compat.py

@@ -14,7 +14,10 @@
 if is_py2:
     from urlparse import urljoin, urlparse
 
-    str = unicode
+    def _str(obj):
+        return unicode(obj)
+
+    str = _str
 
     import io
 

src/code42cli/errors.py

@@ -1,12 +1,24 @@
 ERRORED = False
 
+
 _FORMAT_VALUE_ERROR_MESSAGE = (
     u"input must be a date/time string (e.g. 'YYYY-MM-DD', "
     u"'YY-MM-DD HH:MM', 'YY-MM-DD HH:MM:SS'), or a short value in days, "
     u"hours, or minutes (e.g. 30d, 24h, 15m)"
 )
 
 
+class BadFileError(Exception):
+    def __init__(self, file_path, *args, **kwargs):
+        self.file_path = file_path
+        super(BadFileError, self).__init__()
+
+
+class EmptyFileError(BadFileError):
+    def __init__(self, file_path):
+        super(EmptyFileError, self).__init__(file_path, u"Given empty file {}.".format(file_path))
+
+
 class Code42CLIError(Exception):
     pass
 

src/code42cli/file_readers.py

@@ -0,0 +1,56 @@
+import csv
+
+from code42cli.errors import BadFileError
+
+
+class CliFileReader(object):
+    _ROWS_COUNT = -1
+
+    def __init__(self, file_path):
+        self.file_path = file_path
+
+    def __call__(self, *args, **kwargs):
+        pass
+
+    def get_rows_count(self):
+        if self._ROWS_COUNT == -1:
+            self._ROWS_COUNT = sum(1 for _ in open(self.file_path))
+        if self._ROWS_COUNT == 0:
+            raise BadFileError(u"Given empty file {}.".format(self.file_path))
+        return self._ROWS_COUNT
+
+
+class CSVReader(CliFileReader):
+    """A generator that yields header keys mapped to row values from a csv file."""
+
+    def __init__(self, file_path):
+        with open(file_path) as f:
+            try:
+                self.has_header = csv.Sniffer().has_header(next(f))
+            except StopIteration:
+                raise BadFileError(u"Given empty file {}.".format(file_path))
+        super(CSVReader, self).__init__(file_path)
+
+    def __call__(self, *args, **kwargs):
+        for row in csv.DictReader(kwargs.get(u"bulk_file")):
+            yield row
+
+    def get_rows_count(self):
+        rows_count = super(CSVReader, self).get_rows_count()
+        return rows_count - 1 if self.has_header else rows_count
+
+
+class FlatFileReader(CliFileReader):
+    """A generator that yields a single-value per row from a file."""
+
+    def __call__(self, *args, **kwargs):
+        for row in kwargs[u"bulk_file"]:
+            yield row
+
+
+def create_csv_reader(file_path):
+    return CSVReader(file_path)
+
+
+def create_flat_file_reader(file_path):
+    return FlatFileReader(file_path)

src/code42cli/logger.py

@@ -4,7 +4,7 @@
 import copy
 
 from code42cli.compat import str
-from code42cli.util import get_user_project_path, is_interactive
+from code42cli.util import get_user_project_path, is_interactive, color_text_red
 
 
 logger_deps_lock = Lock()
@@ -124,7 +124,39 @@ def _create_formatter_for_error_file():
 
 
 def _get_red_error_text(text):
-    return u"\033[91mERROR: {}\033[0m".format(text)
+    return color_text_red(u"ERROR: {}".format(text))
+
+
+def get_progress_logger(handler=None):
+    logger = logging.getLogger(u"code42cli_progress_bar")
+    if logger_has_handlers(logger):
+        return logger
+
+    with logger_deps_lock:
+        if not logger_has_handlers(logger):
+            handler = handler or InPlaceStreamHandler()
+            formatter = _get_standard_formatter()
+            logger.setLevel(logging.INFO)
+            return add_handler_to_logger(logger, handler, formatter)
+    return logger
+
+
+class InPlaceStreamHandler(logging.StreamHandler):
+    def __init__(self):
+        super(InPlaceStreamHandler, self).__init__(sys.stdout)
+
+    def emit(self, record):
+        # Borrowed some from python3's logging.StreamHandler to make work on python2.
+        try:
+            msg = u"\r{}\r".format(self.format(record))
+            stream = self.stream
+            stream.write(msg)
+            self.flush()
+        except RuntimeError as err:
+            if u"recursion" in str(err):
+                raise
+        except Exception:
+            self.handleError(record)
 
 
 class CliLogger(object):

src/code42cli/progress_bar.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+
+from code42cli.logger import get_main_cli_logger, get_progress_logger
+
+
+class ProgressBar(object):
+    _FILL = u"█"
+    _LENGTH = 100
+
+    def __init__(self, total_items, logger=None):
+        self._total_items = total_items
+        self._logger = logger or get_progress_logger()
+
+    def update(self, iteration, message):
+        bar = self._create_bar(iteration)
+        progress = u"{}  {} ".format(bar, message)
+        self._logger.info(progress)
+
+    def _create_bar(self, iteration):
+        fill_length = self._calculate_fill_length(iteration)
+        return self._FILL * fill_length + u"-" * (self._LENGTH - fill_length)
+
+    def _calculate_fill_length(self, idx):
+        filled_length = int(self._LENGTH * idx // self._total_items)
+        return filled_length
+
+    def clear_bar_and_print_final(self, final_message):
+        clear = (self._LENGTH + len(final_message)) * u" "
+        self._logger.info(u"{}{}\n".format(final_message, clear))