|
4 | 4 |
|
5 | 5 | from code42cli.compat import open, str |
6 | 6 | from code42cli.worker import Worker |
| 7 | +from code42cli.errors import print_errors_occurred |
| 8 | +from code42cli.args import SDK_ARG_NAME, PROFILE_ARG_NAME |
7 | 9 |
|
8 | 10 |
|
9 | 11 | def generate_template(handler, path=None): |
10 | | - """Looks at the parameter names of `handler` and creates a csv file with the same column names. |
| 12 | + """Looks at the parameter names of `handler` and creates a file with the same column names. If |
| 13 | + `handler` only has one parameter that is not `sdk` or `profile`, it will create a blank file. |
| 14 | + This is useful for commands such as `remove` which only require a list of users. |
11 | 15 | """ |
12 | | - if callable(handler): |
13 | | - argspec = inspect.getargspec(handler) |
14 | | - columns = [str(arg) for arg in argspec.args if arg not in [u"sdk", u"profile"]] |
15 | | - path = path or u"{0}/{1}.csv".format(os.getcwd(), str(handler.__name__)) |
16 | | - _write_template_file(path, columns) |
| 16 | + path = path or u"{0}/{1}.csv".format(os.getcwd(), str(handler.__name__)) |
| 17 | + args = [ |
| 18 | + arg |
| 19 | + for arg in inspect.getargspec(handler).args |
| 20 | + if arg != SDK_ARG_NAME and arg != PROFILE_ARG_NAME |
| 21 | + ] |
17 | 22 |
|
| 23 | + if len(args) <= 1: |
| 24 | + print( |
| 25 | + u"A blank file was generated because there are no csv headers needed for this command. " |
| 26 | + u"Simply enter one {} per line.".format(args[0]) |
| 27 | + ) |
| 28 | + # Set args to None so that we don't make a header out of the single arg. |
| 29 | + args = None |
18 | 30 |
|
19 | | -def _write_template_file(path, columns): |
20 | | - with open(path, u"w", encoding=u"utf8") as new_csv: |
21 | | - new_csv.write(u",".join(columns)) |
| 31 | + _write_template_file(path, args) |
22 | 32 |
|
23 | 33 |
|
24 | | -def run_bulk_process(csv_file_path, row_handler): |
25 | | - processor = _create_bulk_processor(csv_file_path, row_handler) |
| 34 | +def _write_template_file(path, columns=None): |
| 35 | + with open(path, u"w", encoding=u"utf8") as new_file: |
| 36 | + if columns: |
| 37 | + new_file.write(u",".join(columns)) |
| 38 | + |
| 39 | + |
| 40 | +def run_bulk_process(file_path, row_handler, reader=None): |
| 41 | + """Runs a bulk process. |
| 42 | + |
| 43 | + Args: |
| 44 | + file_path (str or unicode): The path to the file feeding the data for the bulk process. |
| 45 | + row_handler (callable): A callable that you define to process values from the row as |
| 46 | + either *args or **kwargs. |
| 47 | + reader: (CSVReader or FlatFileReader, optional): A generator that reads rows and yields data into |
| 48 | + `row_handler`. If None, it will use a CSVReader. Defaults to None. |
| 49 | + """ |
| 50 | + reader = reader or CSVReader() |
| 51 | + processor = _create_bulk_processor(file_path, row_handler, reader) |
26 | 52 | processor.run() |
27 | 53 |
|
28 | 54 |
|
29 | | -def _create_bulk_processor(csv_file_path, row_handler): |
| 55 | +def _create_bulk_processor(file_path, row_handler, reader): |
30 | 56 | """A factory method to create the bulk processor, useful for testing purposes.""" |
31 | | - return BulkProcessor(csv_file_path, row_handler) |
| 57 | + return BulkProcessor(file_path, row_handler, reader) |
32 | 58 |
|
33 | 59 |
|
34 | 60 | class BulkProcessor(object): |
35 | | - """A class for bulk processing a csv file. |
| 61 | + """A class for bulk processing a file. |
36 | 62 | |
37 | 63 | Args: |
38 | | - csv_file_path (str or unicode): The path to the csv file for processing. |
39 | | - row_handler (callable): To be executed on each row given **kwargs representing the column |
40 | | - names mapped to the properties found in the row. For example, if the csv file header |
41 | | - looked like `prop_a,prop_b` and the next row looked like `1,test`, then row handler |
42 | | - would receive args `prop_a: '1', prop_b: 'test'` when processing row 1. |
| 64 | + file_path (str or unicode): The path to the file for processing. |
| 65 | + row_handler (callable): A callable that you define to process values from the row as |
| 66 | + either *args or **kwargs. For example, if it's a csv file with header `prop_a,prop_b` |
| 67 | + and first row `1,test`, then `row_handler` should receive kwargs |
| 68 | + `prop_a: '1', prop_b: 'test'` when processing the first row. If it's a flat file, then |
| 69 | + `row_handler` only needs to take an extra arg. |
| 70 | + reader (CSVReader or FlatFileReader): A generator that reads rows and yields data into `row_handler`. |
43 | 71 | """ |
44 | 72 |
|
45 | | - def __init__(self, csv_file_path, row_handler): |
46 | | - self.csv_file_path = csv_file_path |
| 73 | + def __init__(self, file_path, row_handler, reader): |
| 74 | + self.file_path = file_path |
47 | 75 | self._row_handler = row_handler |
| 76 | + self._reader = reader |
48 | 77 | self.__worker = Worker(5) |
49 | 78 |
|
50 | 79 | def run(self): |
51 | 80 | """Processes the csv file specified in the ctor, calling `self.row_handler` on each row.""" |
52 | | - with open(self.csv_file_path, newline=u"", encoding=u"utf8") as csv_file: |
53 | | - rows = _create_dict_reader(csv_file) |
54 | | - self._process_rows(rows) |
| 81 | + with open(self.file_path, newline=u"", encoding=u"utf8") as bulk_file: |
| 82 | + for row in self._reader(bulk_file=bulk_file): |
| 83 | + self._process_row(row) |
55 | 84 | self.__worker.wait() |
| 85 | + self._print_result() |
| 86 | + |
| 87 | + def _process_row(self, row): |
| 88 | + if type(row) is dict: |
| 89 | + self._process_csv_row(row) |
| 90 | + elif row: |
| 91 | + self._process_flat_file_row(row.strip()) |
| 92 | + |
| 93 | + def _process_csv_row(self, row): |
| 94 | + # Removes problems from including extra comments. Error messages from out of order args |
| 95 | + # are more indicative this way too. |
| 96 | + row.pop(None, None) |
| 97 | + self.__worker.do_async(lambda *args, **kwargs: self._row_handler(*args, **kwargs), **row) |
| 98 | + |
| 99 | + def _process_flat_file_row(self, row): |
| 100 | + if row: |
| 101 | + self.__worker.do_async(lambda *args, **kwargs: self._row_handler(*args, **kwargs), row) |
| 102 | + |
| 103 | + def _print_result(self): |
| 104 | + stats = self.__worker.stats |
| 105 | + successes = stats.total - stats.total_errors |
| 106 | + print(u"{} processed successfully out of {}.".format(successes, stats.total)) |
| 107 | + if stats.total_errors: |
| 108 | + print_errors_occurred() |
| 109 | + |
| 110 | + |
| 111 | +class CSVReader(object): |
| 112 | + """A generator that yields header keys mapped to row values from a csv file.""" |
| 113 | + |
| 114 | + def __call__(self, *args, **kwargs): |
| 115 | + for row in csv.DictReader(kwargs.get(u"bulk_file")): |
| 116 | + yield row |
56 | 117 |
|
57 | | - def _process_rows(self, rows): |
58 | | - for row in rows: |
59 | | - self.__worker.do_async(lambda **kwargs: self._row_handler(**kwargs), **row) |
60 | 118 |
|
| 119 | +class FlatFileReader(object): |
| 120 | + """A generator that yields a single-value per row from a file.""" |
61 | 121 |
|
62 | | -def _create_dict_reader(csv_file): |
63 | | - return csv.DictReader(csv_file) |
| 122 | + def __call__(self, *args, **kwargs): |
| 123 | + for row in kwargs[u"bulk_file"]: |
| 124 | + yield row |
0 commit comments