improved code structure

gcode-tools/gcode-shifter/gcode_shift.py

@@ -5,6 +5,7 @@ import os
 import re
 import sys
 import math
+import collections
 from decimal import Decimal as d
 
 AXIS = "XYZ"
@@ -13,13 +14,16 @@ AXIS_REGEX = r"\b([%s][\+\-0-9\.]+)\b" % AXIS
 
 class GCodeFilter(object):
 
-    def __init__(self, lower, upper, func=None):
+    def __init__(self, destination, lower, upper, func=None):
         self.pos = tuple([d(0)] * len(AXIS))
         self.lower = lower
         self.upper = upper
         self.recalc_func = func
         self.source_pos = PositionHandler()
         self.target_pos = PositionHandler()
+        self.destination = destination
+        self.destination_buffer = collections.deque()
+        self._status = {}
 
     def is_inside(self, position):
         for i in range(len(position)):
@@ -45,7 +49,7 @@ class GCodeFilter(object):
         else:
             return False
 
-    def get_processed_lines(self, line):
+    def parse(self, line):
         """ parse the line, update current position and replace line string
         """
         if self.update_position(line):
@@ -53,19 +57,85 @@ class GCodeFilter(object):
             handler = LineHandler(line)
             positions = list(self.transform_position())
             if not positions:
-                yield handler.get_line(None, None)
+                self.destination_buffer.append(handler.get_line(None, None))
             else:
                 for pos in positions:
                     if callable(pos):
                         pos, line = pos()
-                        yield line
+                        self.destination_buffer.append(line)
                     else:
                         changed_axes = self.target_pos.get_changed_axes(pos)
-                        yield handler.get_line(pos, changed_axes)
+                        self.destination_buffer.append(handler.get_line(pos,
+                                changed_axes))
                     self.target_pos.update(pos)
         else:
             # no coordinate given
-            yield line
+            self.destination_buffer.append(line)
+        self.write()
+
+    def _parse_rapid(self, line):
+        # return True / False / None for G0/G1/other commands
+        for item in line.upper().split():
+            if item == "G0":
+                return True
+            if item == "G1":
+                return False
+        return None
+
+    def _parse_rapid_noop(self, line):
+        rapid = self._parse_rapid(line)
+        if rapid is None:
+            return False
+        else:
+            # only one item (G0 or G1)
+            return len(line.strip().split()) == 1
+
+    def write(self):
+        """ add all no-op filters here (e.g. "GO / G1" -> "G1")
+        """
+        if not self.destination:
+            raise IOError("GCodeFilter was closed before")
+        pending_rapid = None
+        while self.destination_buffer:
+            item = self.destination_buffer.popleft()
+            if item is None:
+                # skip invalid items
+                continue
+            rapid = self._parse_rapid(item)
+            rapid_noop = self._parse_rapid_noop(item)
+            if rapid_noop:
+                if rapid == self._status.get("rapid"):
+                    # skip this no-op (and any pending ones)
+                    pending_rapid = None
+                    continue
+                else:
+                    pending_rapid = item
+                    continue
+            else:
+                if (not pending_rapid is None) and (rapid is None) and \
+                        (self._parse_rapid(pending_rapid) != rapid):
+                    self.destination.write(pending_rapid)
+                pending_rapid = None
+            # no matches -> write it
+            self.destination.write(item)
+            if not rapid is None:
+                self._status["rapid"] = rapid
+        if not pending_rapid is None:
+            # put it back to the list
+            self.destination_buffer.appendleft(pending_rapid)
+
+    def close(self):
+        if self.destination is None:
+            return
+        self.write()
+        # write all pending items
+        while self.destination_buffer:
+            self.destination.write(self.destination_buffer.popleft())
+        self.destination.close()
+        self.destination = None
+
+    def __del__(self):
+        self.close()
 
 
 class PositionHandler(object):
@@ -355,24 +425,23 @@ if __name__ == "__main__":
     options = parser.parse_args()
     low = (options.minx, options.miny, options.minz)
     high = (options.maxx, options.maxy, options.maxz)
+    common_args = (outfile, low, high)
     if options.action == "shift":
         shift = (options.shiftx, options.shifty, options.shiftz)
-        gcode_filter = ShiftFilter(low, high, shift=shift)
+        gcode_filter = ShiftFilter(*common_args, shift=shift)
     elif options.action == "crop":
-        gcode_filter = CropFilter(low, high)
+        gcode_filter = CropFilter(*common_args)
     elif options.action == "rotate":
         matrix = get_rotate_matrix(options.rotate_axis, options.rotate_angle)
-        gcode_filter = MatrixFilter(low, high, matrix=matrix)
+        gcode_filter = MatrixFilter(*common_args, matrix=matrix)
     elif options.action == "densify":
-        gcode_filter = DensifyFilter(low, high, densify_dir=options.densify_dir,
+        gcode_filter = DensifyFilter(*common_args,
+                densify_dir=options.densify_dir,
                 densify_loops=options.densify_loops,
                 densify_minimum_step=options.densify_minimum_step)
     else:
         print >>sys.stderr, "No valid action choosen"
         sys.exit(1)
     for line in infile.readlines():
-        for out_line in gcode_filter.get_processed_lines(line):
+        gcode_filter.parse(line)
-            # omit empty lines
-            if not out_line is None:
-                outfile.write(out_line)