|
26 | 26 | import os |
27 | 27 | import pathlib |
28 | 28 | import sys |
| 29 | +from contextlib import nullcontext |
29 | 30 |
|
30 | 31 | from progress.spinner import Spinner |
31 | 32 |
|
@@ -105,48 +106,46 @@ def count_files(self, scan_dir: str) -> bool: |
105 | 106 | """ |
106 | 107 | success = True |
107 | 108 | if not scan_dir: |
108 | | - raise Exception(f'ERROR: Please specify a folder to scan') |
| 109 | + raise Exception('ERROR: Please specify a folder to scan') |
109 | 110 | if not os.path.exists(scan_dir) or not os.path.isdir(scan_dir): |
110 | 111 | raise Exception(f'ERROR: Specified folder does not exist or is not a folder: {scan_dir}') |
111 | 112 |
|
112 | 113 | self.print_msg(f'Searching {scan_dir} for files to count...') |
113 | | - spinner = None |
114 | | - if not self.quiet and self.isatty: |
115 | | - spinner = Spinner('Searching ') |
116 | | - file_types = {} |
117 | | - file_count = 0 |
118 | | - file_size = 0 |
119 | | - for root, dirs, files in os.walk(scan_dir): |
120 | | - self.print_trace(f'U Root: {root}, Dirs: {dirs}, Files {files}') |
121 | | - dirs[:] = self.__filter_dirs(dirs) # Strip out unwanted directories |
122 | | - filtered_files = self.__filter_files(files) # Strip out unwanted files |
123 | | - self.print_trace(f'F Root: {root}, Dirs: {dirs}, Files {filtered_files}') |
124 | | - for file in filtered_files: # Cycle through each filtered file |
125 | | - path = os.path.join(root, file) |
126 | | - f_size = 0 |
127 | | - try: |
128 | | - f_size = os.stat(path).st_size |
129 | | - except Exception as e: |
130 | | - self.print_trace(f'Ignoring missing symlink file: {file} ({e})') # broken symlink |
131 | | - if f_size > 0: # Ignore broken links and empty files |
132 | | - file_count = file_count + 1 |
133 | | - file_size = file_size + f_size |
134 | | - f_suffix = pathlib.Path(file).suffix |
135 | | - if not f_suffix or f_suffix == '': |
136 | | - f_suffix = 'no_suffix' |
137 | | - self.print_trace(f'Counting {path} ({f_suffix} - {f_size})..') |
138 | | - fc = file_types.get(f_suffix) |
139 | | - if not fc: |
140 | | - fc = [1, f_size] |
141 | | - else: |
142 | | - fc[0] = fc[0] + 1 |
143 | | - fc[1] = fc[1] + f_size |
144 | | - file_types[f_suffix] = fc |
145 | | - if spinner: |
146 | | - spinner.next() |
147 | | - # End for loop |
148 | | - if spinner: |
149 | | - spinner.finish() |
| 114 | + spinner_ctx = Spinner('Searching ') if (not self.quiet and self.isatty) else nullcontext() |
| 115 | + |
| 116 | + with spinner_ctx as spinner: |
| 117 | + file_types = {} |
| 118 | + file_count = 0 |
| 119 | + file_size = 0 |
| 120 | + for root, dirs, files in os.walk(scan_dir): |
| 121 | + self.print_trace(f'U Root: {root}, Dirs: {dirs}, Files {files}') |
| 122 | + dirs[:] = self.__filter_dirs(dirs) # Strip out unwanted directories |
| 123 | + filtered_files = self.__filter_files(files) # Strip out unwanted files |
| 124 | + self.print_trace(f'F Root: {root}, Dirs: {dirs}, Files {filtered_files}') |
| 125 | + for file in filtered_files: # Cycle through each filtered file |
| 126 | + path = os.path.join(root, file) |
| 127 | + f_size = 0 |
| 128 | + try: |
| 129 | + f_size = os.stat(path).st_size |
| 130 | + except Exception as e: |
| 131 | + self.print_trace(f'Ignoring missing symlink file: {file} ({e})') # broken symlink |
| 132 | + if f_size > 0: # Ignore broken links and empty files |
| 133 | + file_count = file_count + 1 |
| 134 | + file_size = file_size + f_size |
| 135 | + f_suffix = pathlib.Path(file).suffix |
| 136 | + if not f_suffix or f_suffix == '': |
| 137 | + f_suffix = 'no_suffix' |
| 138 | + self.print_trace(f'Counting {path} ({f_suffix} - {f_size})..') |
| 139 | + fc = file_types.get(f_suffix) |
| 140 | + if not fc: |
| 141 | + fc = [1, f_size] |
| 142 | + else: |
| 143 | + fc[0] = fc[0] + 1 |
| 144 | + fc[1] = fc[1] + f_size |
| 145 | + file_types[f_suffix] = fc |
| 146 | + if spinner: |
| 147 | + spinner.next() |
| 148 | + # End for loop |
150 | 149 | self.print_stderr(f'Found {file_count:,.0f} files with a total size of {file_size / (1 << 20):,.2f} MB.') |
151 | 150 | if file_types: |
152 | 151 | csv_dict = [] |
|
0 commit comments