Last active
August 9, 2019 19:34
-
-
Save hughdbrown/86a1973ae5b58e4ac82b44e128b04f38 to your computer and use it in GitHub Desktop.
All your absolute_importing needs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
import os | |
class SedFile(object): | |
def __init__(self, filename): | |
self.filename = filename | |
self.modified = False | |
with open(self.filename, "r") as handle: | |
self.data = [line.rstrip() for line in handle] | |
def __enter__(self): | |
return self | |
def __exit__(self, type, value, tb, ): | |
if self.modified: | |
with open(self.filename, "w") as handle: | |
handle.write('\n'.join(self.data) + '\n') | |
def modify(self): | |
raise NotImplementedError | |
class AbsoluteImportInjector(SedFile): | |
""" | |
Derivative of SedFile dedicated to injecting | |
from __future__ import absolute_import | |
into python files. | |
""" | |
prefix = "from __future__ import " | |
injection = 'from __future__ import absolute_import\n' | |
def modify(self): | |
self._modify_absolute() | |
self._modify_multiple() | |
def _modify_absolute(self): | |
""" | |
Inject absolute import into file if not already present. | |
Most of the logic is dedicated to findint the correct injection point | |
""" | |
if any(self.data) and not any("absolute_import" in line for line in self.data): | |
print(self.filename) | |
i = 0 | |
if self.data[i].startswith(('#!/')): | |
# Skip execution header | |
i += 1 | |
if self.data[i].startswith('#') and 'coding' in self.data[i]: | |
# Skip file encoding comment | |
i += 1 | |
comment_header = ("'''", '"""') | |
if self.data[i].lstrip().startswith(comment_header): | |
# Skip any docstring | |
if len(self.data[i]) == 3 or not self.data[i].endswith(comment_header): | |
# If the comment does not start and end on this line ... | |
i += 1 | |
try: | |
while not self.data[i].endswith(comment_header): | |
print(i, self.data[i]) | |
i += 1 | |
except IndexError: | |
return None | |
assert self.data[i].endswith(comment_header) | |
i += 1 | |
for j, line in enumerate(self.data, start=i): | |
# Skip other comments | |
if not line.lstrip().startswith('#'): | |
break | |
i = j | |
# This is the first line that a future import might be placed | |
if self.data[i].startswith(self.prefix): | |
self.data[i] += ", absolute_import" | |
else: | |
self.data = self.data[:i] + [self.injection] + self.data[i:] | |
self.modified = True | |
def _modify_multiple(self): | |
"""Flatten multiple (possibly non-conecutive) matches into a single line""" | |
x = [i for i, line in enumerate(self.data) if line.startswith(self.prefix)] | |
if len(x) > 1: | |
# Join lines onto the first import-line, preserving order | |
joined_lines = x[1:] | |
for j in joined_lines: | |
absolute_extract = self.data[j][len(self.prefix):] | |
self.data[x[0]] += ", {}".format(absolute_extract) | |
# Discard lines that have been joined | |
self.data = [line for j, line in enumerate(self.data) if j not in set(joined_lines)] | |
self.modified = True | |
def path_iter(dir='.', exts=('.py', )): | |
""" | |
Generator for files in tree that match any exts | |
""" | |
for root, _, files in os.walk(dir): | |
for filename in files: | |
fullpath = os.path.join(os.path.normpath(root), filename) | |
if os.path.splitext(fullpath)[1] in exts: | |
yield fullpath | |
def main(): | |
""" | |
Traverse down tree, looking for files to modify. Save modified files | |
""" | |
for fullpath in path_iter(): | |
with AbsoluteImportInjector(fullpath) as obj: | |
obj.modify() | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment