def _find_files(paths, extensions):
if paths is None:
return
for path in paths:
if os.path.isfile(path):
yield path
else:
for root, dirs, files in os.walk(path):
for f in files:
ext = os.path.splitext(f)[-1].lstrip('.')
if extensions is None:
log.error('must pass --extension when --input is a directory')
sys.exit(1)
if ext in extensions:
yield os.path.join(root, f)
for d in dirs[:]:
if d != "." * len(d) and d.startswith("."): # ignore .*
dirs.remove(d)
def main():
"""Handle and process arguments from sys.argv."""
logger.setup()
args = _handle_arguments()
logger.setup(args.verbose) # Reset logging level
if args.multiprocess <= 0:
log.error('number of processes must be > 0')
sys.exit(1)
processor = _file_processor(args.pattern, args.dry_run)
files = list(_find_files(args.input, args.extension))
if bool(files) != bool(args.input):
log.error('could not find any files for the given paths and extension')
sys.exit(1)
if not files: # Single process mode if stdin
log.info('running in stdin/stdout mode')
processor(None)
elif len(files) == 1 or args.multiprocess == 1: # Single process if only one file or only one process
log.info('running across {} file(s) using a single process'
.format(len(files)))
processor(files[0])
else:
process_pool = multiprocessing.Pool(args.multiprocess)
try:
result = process_pool.map_async(
processor,
files,
)
process_pool.close()
log.info('running across {} file(s) using {} processes'
.format(len(files), args.multiprocess))
# Cannot do process_pool.wait() because it prevents KeyboardInterrupt from being sent
# See http://stackoverflow.com/questions/1408356/keyboard-interrupts-with-pythons-multiprocessing-pool
while not result.ready():
time.sleep(0.01)
if not result.successful():
log.error('multiprocessing failed (are your replace functions pickleable?)')
sys.exit(1)
result = result.get()
assert len(result) == len(files)
if not all(result):
log.error('failed to process {} files'
.format(len(result) - sum(result)))
sys.exit(1)
except:
process_pool.terminate()
raise
finally:
process_pool.join()
def _get_patterns(pattern):
if hasattr(pattern, "patterns"):
return pattern.patterns
if hasattr(pattern, "grammar") and hasattr(pattern, "replace"):
patterns = [(pattern.grammar, pattern.replace)]
if hasattr(pattern, "extra"):
patterns.append(get_pattern_for_extra(pattern.extra))
return patterns
log.error(
'pattern file {} must define either ("patterns") or ("grammar" and "replace" and optionally "extra")'
.format(pattern.__name__))
sys.exit(1)
def get_patterns(*pattern_modules):
"""Returns patterns for pattern modules."""
pattern_list = []
for pattern in pattern_modules:
# new style pattern files
if hasattr(pattern, "patterns"):
patterns = pattern.patterns
# old style pattern files
elif hasattr(pattern, "grammar") and hasattr(pattern, "replace"):
patterns = [(pattern.grammar, pattern.replace)]
if hasattr(pattern, "extra"):
patterns.append(get_pattern_for_extra(pattern.extra))
else:
log.error(
'pattern file {} must define either ("patterns") or ("grammar" and "replace" and optionally "extra")'
.format(pattern.__name__)
)
sys.exit(1)
pattern_list.append(patterns)
return pattern_list