jklemm · October 9, 2019 01:39 · Oct 9, 2019 · May 15, 2019
diff --git a/adiciona_unicode_literals.py b/adiciona_unicode_literals.py
@@ -53,24 +53,23 @@ def add_unicode_literals(one_file):
 
         original_file.seek(0)
         lines = conteudo.splitlines()
+
+        adicionar = 0
 
         if 'coding' in conteudo:
-            lines[0] = '# coding: utf-8'
-            lines.insert(1, UNICODE_LITERALS)
-            if lines[2] != '':
-                lines.insert(2, '')
-            if 'class' in lines[3]:
-                lines.insert(3, '')
-            if 'def' in lines[3]:
-                lines.insert(3, '')
-        else:
-            lines.insert(0, UNICODE_LITERALS)
-            if lines[1] != '':
-                lines.insert(1, '')
-            if 'class' in lines[2]:
-                lines.insert(2, '')
-            if 'def' in lines[2]:
-                lines.insert(2, '')
+            adicionar += 1
+        if 'absolute_import' in conteudo:
+            adicionar += 1
+        if 'print_function' in conteudo:
+            adicionar += 1
+
+        lines.insert(adicionar, UNICODE_LITERALS)
+        if lines[adicionar + 1] != '':
+            lines.insert(adicionar + 1, '')
+        if 'class' in lines[adicionar + 2]:
+            lines.insert(adicionar + 2, '')
+        if 'def' in lines[adicionar + 2]:
+            lines.insert(adicionar + 2, '')
 
     lines.append('')
 

diff --git a/adiciona_unicode_literals.py b/adiciona_unicode_literals.py
@@ -0,0 +1,104 @@
+#!/usr/bin/python
+# coding: utf-8
+from __future__ import absolute_import
+
+import sys
+import os
+from builtins import input
+
+UNICODE_LITERALS = 'from __future__ import unicode_literals'
+
+ALLOWED_EXTENSION = '.py'
+IGNORED_DIRS = ('__init__', '.idea', '.sass-cache', 'media', 'migrations', 'node_modules', 'static', 'templates')
+
+
+def listar_arquivos_sem_unicode_literals(directory):
+    file_list = []
+    for root, subdirs, files in os.walk(directory):
+        for filename in files:
+
+            if not filename.endswith(ALLOWED_EXTENSION):
+                continue
+
+            file_path = os.path.join(root, filename)
+
+            if any([word in file_path for word in IGNORED_DIRS]):
+                continue
+
+            valid_file = is_file_valid(file_path)
+            if not valid_file:
+                continue
+
+            have_import = is_file_have_unicode_literals(file_path)
+            if not have_import:
+                file_list.append(file_path)
+    return file_list
+
+
+def is_file_valid(one_file):
+    with open(one_file, 'r') as myfile:
+        num_lines = sum(1 for line in myfile)
+    return num_lines >= 4
+
+
+def is_file_have_unicode_literals(one_file):
+    with open(one_file, 'r') as myfile:
+        content = myfile.read()
+        return UNICODE_LITERALS in content
+
+
+def add_unicode_literals(one_file):
+    with open(one_file, 'r') as original_file:
+        conteudo = original_file.read()
+
+        original_file.seek(0)
+        lines = conteudo.splitlines()
+
+        if 'coding' in conteudo:
+            lines[0] = '# coding: utf-8'
+            lines.insert(1, UNICODE_LITERALS)
+            if lines[2] != '':
+                lines.insert(2, '')
+            if 'class' in lines[3]:
+                lines.insert(3, '')
+            if 'def' in lines[3]:
+                lines.insert(3, '')
+        else:
+            lines.insert(0, UNICODE_LITERALS)
+            if lines[1] != '':
+                lines.insert(1, '')
+            if 'class' in lines[2]:
+                lines.insert(2, '')
+            if 'def' in lines[2]:
+                lines.insert(2, '')
+
+    lines.append('')
+
+    with open(one_file, 'w') as output_file:
+        new_file_content = '\n'.join(lines)
+        output_file.write(new_file_content)
+
+
+if __name__ == '__main__':
+    if len(sys.argv) == 1:
+        print('Error, you need to inform a directory')
+        exit(1)
+
+    diretorio = sys.argv[1]
+
+    lista_de_arquivos = listar_arquivos_sem_unicode_literals(diretorio)
+
+    quantidade = len(lista_de_arquivos)
+
+    if quantidade > 200:
+        resposta = input('Um total de {} arquivos serão alterados, deseja prosseguir? '.format(quantidade))
+
+        if resposta.upper() != 'S':
+            exit(0)
+
+    for arquivo in lista_de_arquivos:
+        print('- {}'.format(arquivo))
+        add_unicode_literals(arquivo)
+    print('{} arquivos foram alterados.'.format(quantidade))
+
+    exit(0)