tidy: clenaup and refactoring

yanne · yanne · commit 9ee0f07b00db · 2020-01-09T11:43:17.000+02:00
- add `save()` method to model
- simplify code by inlining `datafilewriter`
diff --git a/src/robot/parsing/builders.py b/src/robot/parsing/builders.py
@@ -23,16 +23,16 @@
 
 def get_model(source, data_only=False, curdir=None):
     tokens = get_tokens(source, data_only)
-    return _build_model(get_statements(tokens, curdir))
+    return _build_model(source, get_statements(tokens, curdir))
 
 
 def get_resource_model(source, data_only=False, curdir=None):
     tokens = get_resource_tokens(source, data_only)
-    return _build_model(get_statements(tokens, curdir))
+    return _build_model(source, get_statements(tokens, curdir))
 
 
-def _build_model(statements):
-    builder = FileBuilder()
+def _build_model(source, statements):
+    builder = FileBuilder(source)
     stack = [builder]
     for statement in statements:
         while not stack[-1].handles(statement):
@@ -57,8 +57,8 @@ def statement(self, statement):
 
 class FileBuilder(Builder):
 
-    def __init__(self, model=None):
-        Builder.__init__(self, model or File())
+    def __init__(self, source=None):
+        Builder.__init__(self, File(source))
 
     def statement(self, statement):
         try:
diff --git a/src/robot/parsing/model/blocks.py b/src/robot/parsing/model/blocks.py
@@ -15,25 +15,45 @@
 
 import ast
 
-from robot.utils import normalize_whitespace
+from robot.utils import normalize_whitespace, file_writer
 
+from .visitor import ModelVisitor
 from ..lexer import Token
 
 
+class FileWriter(ModelVisitor):
+
+    def __init__(self, output):
+        self.output = output
+
+    def visit_Statement(self, statement):
+        for line in statement.lines:
+            for token in line:
+                self.output.write(token.value)
+
+
 class Block(ast.AST):
     _fields = ()
 
 
 class File(Block):
     _fields = ('sections',)
 
-    def __init__(self):
+    def __init__(self, source=None):
+        self.source = source
         self.sections = []
 
     @property
     def has_tests(self):
         return any(isinstance(s, TestCaseSection) for s in self.sections)
 
+    def save(self, given_output, newline=None):
+        output = given_output if given_output else \
+            file_writer(self.source, newline=newline)
+        FileWriter(output).visit(self)
+        if given_output is None:
+            output.close()
+
 
 class Section(Block):
     _fields = ('header', 'body')
diff --git a/src/robot/tidy.py b/src/robot/tidy.py
@@ -39,7 +39,8 @@
 
 from robot.errors import DataError
 from robot.parsing import get_model, SuiteStructureBuilder, SuiteStructureVisitor
-from robot.tidypkg import DataFileWriter
+from robot.tidypkg import (Aligner, Cleaner, SeparatorCleaner, PipeAdder,
+                           NewlineAdder)
 from robot.utils import Application, file_writer
 
 # TODO: maybe rename --format to --extension
@@ -180,8 +181,15 @@ def directory(self, path):
         data = SuiteStructureBuilder().build([path])
         data.visit(self)
 
-    def _tidy(self, ast, output=None):
-        DataFileWriter(output=output, **self._options).write(ast)
+    def _tidy(self, model, output=None):
+        Cleaner().visit(model)
+        NewlineAdder().visit(model)
+        if self._options['pipe_separated']:
+            PipeAdder().visit(model)
+        else:
+            SeparatorCleaner(self._options['txt_separating_spaces']).visit(model)
+        Aligner().visit(model)
+        model.save(output, self._options['line_separator'])
 
     def visit_file(self, file):
         self.inplace(file.source)
diff --git a/src/robot/tidypkg/__init__.py b/src/robot/tidypkg/__init__.py
@@ -13,4 +13,5 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from .datafilewriter import DataFileWriter
+from .transformers import (Aligner, Cleaner, SeparatorCleaner, PipeAdder,
+                           NewlineAdder)
diff --git a/src/robot/tidypkg/datafilewriter.py b/src/robot/tidypkg/datafilewriter.py
diff --git a/src/robot/tidypkg/transformers.py b/src/robot/tidypkg/transformers.py
@@ -111,9 +111,9 @@ def visit_Keyword(self, node):
 
 class SeparatorCleaner(ModelTransformer):
 
-    def __init__(self, ctx):
+    def __init__(self, separator_width):
         self.indent = 0
-        self.separator = ' ' * ctx.txt_separating_spaces
+        self.separator = ' ' * separator_width
 
     def visit_TestCase(self, node):
         self.indent += 1
@@ -155,8 +155,8 @@ def normalize_separator(self, index, token):
 
 class ColumnAligner(ModelTransformer):
 
-    def __init__(self, ctx, widths):
-        self.ctx = ctx
+    def __init__(self, short_test_name_length, widths):
+        self.short_test_name_length = short_test_name_length
         self.widths = widths
         self.test_name_len = 0
         self.indent = 0
@@ -195,8 +195,7 @@ def align_statement(self, statement):
             for token, width in zip(line, widths):
                 exp_pos += width
                 if self.should_write_content_after_name(line_pos):
-                    # FIXME, this hack ony works with indent = 4
-                    exp_pos -= (self.test_name_len - 4)
+                    exp_pos -= self.test_name_len
                     self.first_statement_after_name_seen = True
                 token.value = (exp_pos - line_pos) * ' ' + token.value
                 line_pos += len(token.value)
@@ -213,7 +212,7 @@ def _should_be_indented(self, line):
 
     def should_write_content_after_name(self, line_pos):
         return line_pos == 0 and not self.first_statement_after_name_seen \
-               and self.test_name_len < self.ctx.short_test_name_length
+               and self.test_name_len < self.short_test_name_length
 
 
 class ColumnWidthCounter(ModelTransformer):
@@ -239,15 +238,15 @@ def _count_widths_from_statement(self, statement, indent=0):
 
 
 class Aligner(ModelTransformer):
-
-    def __init__(self, ctx):
-        self.ctx = ctx
+    setting_and_variable_name_length = 14
+    short_test_name_length = 18
 
     def visit_TestCaseSection(self, section):
         if len(section.header.data_tokens) > 1:
             counter = ColumnWidthCounter()
             counter.visit(section)
-            ColumnAligner(self.ctx, counter.widths).visit(section)
+            ColumnAligner(self.short_test_name_length,
+                          counter.widths).visit(section)
         return section
 
     def visit_KeywordSection(self, section):
@@ -262,7 +261,7 @@ def visit_Statement(self, statement):
             first_value = [t for t in line if t.type != Token.SEPARATOR][0]
             value_index = line.index(first_value)
             line[value_index].value = line[value_index].value.ljust(
-                self.ctx.setting_and_variable_name_length)
+                self.setting_and_variable_name_length)
         return statement
 
 
@@ -339,39 +338,3 @@ def visit_ForLoop(self, loop):
                                    loop.end.tokens[0], Token(Token.EOL, '\n'))
         self.generic_visit(loop)
         return loop
-
-
-class Formatter(ModelTransformer):
-
-    def __init__(self, ctx, row_writer):
-        self.ctx = ctx
-        self.writer = row_writer
-        self.indent = 0
-
-    def visit_Section(self, section):
-        self.generic_visit(section)
-
-    def visit_TestCase(self, node):
-        self._write_name(node.name_tokens)
-        self.generic_visit(node.body)
-
-    def visit_Keyword(self, node):
-        self._write_name(node.name_tokens)
-        self.generic_visit(node.body)
-
-    def _write_name(self, name):
-        self._write_statement(name)
-
-    def visit_ForLoop(self, node):
-        self._write_statement(node.header)
-        self.generic_visit(node.body)
-        if node.end:
-            self._write_statement(node.end)
-
-    def visit_Statement(self, statement):
-        self._write_statement(statement)
-
-    def _write_statement(self, statement):
-        for line in statement.lines:
-            self.writer.write(line)
-
