Put non-element classes, and unrecognized code other than imports, in section 3

diff --git a/ast_to_pitree.py b/ast_to_pitree.py
index 32bf5b8..b05abed 100755 (executable)
--- a/ast_to_pitree.py
+++ b/ast_to_pitree.py
@@ -65,15 +65,17 @@ re_comment = re.compile(
 )
 stack = []
 classes = []
-base_classes = [{'element.Element': []}] # params
+base_classes = [set(['element.Element'])]
 method = {}
 
 sys.stdout.write('%{\n')
-had_blank = True
 
-lines = []
+section2 = []
 delayed_line = ''
 
+section3 = []
+had_blank = True
+
 line = sys.stdin.readline()
 while len(line):
   match = re_class.match(line)
@@ -93,10 +95,10 @@ while len(line):
       #  ] = temp_fields
       temp_indent, _, _, _ = stack.pop()
       if len(delayed_line):
-        lines.append('{0:s};\n'.format(delayed_line[:-3]))
+        section2.append('{0:s};\n'.format(delayed_line[:-3]))
         delayed_line = ''
       else:
-        lines.append('{0:s}}};\n'.format(temp_indent))
+        section2.append('{0:s}}};\n'.format(temp_indent))
     for i in range(len(base_classes) - 1, -1, -1):
       if base_class in base_classes[i]:
         classes.append(
@@ -105,19 +107,35 @@ while len(line):
         full_base_class = '.'.join(
           [j for _, j, _, _ in stack[:i]] + [base_class]
         )
-        base_classes[-1][class_name] = list(base_classes[i][base_class])
+        base_classes[-1].add(class_name)
+        stack.append((indent, class_name, base_class, full_base_class))
+        base_classes.append(set())
+        section2.append(delayed_line)
+        delayed_line = '{0:s}class {1:s}{2:s} {{\n'.format(
+          indent,
+          class_name,
+          '' if base_class == 'element.Element' else ': ' + base_class
+        )
+        line = sys.stdin.readline()
         break
     else:
-      full_base_class = base_class
-    stack.append((indent, class_name, base_class, full_base_class))
-    base_classes.append({})
-    lines.append(delayed_line)
-    delayed_line = '{0:s}class {1:s}{2:s} {{\n'.format(
-      indent,
-      class_name,
-      ': ' + base_class if len(base_class) else ''
-    )
-    line = sys.stdin.readline()
+      # unrecognized class, copy to header part directly
+      line_rstrip = line.rstrip()
+      while True:
+        if len(line_rstrip):
+          section3.append(line_rstrip + '\n')
+          had_blank = False
+        elif not had_blank:
+          section3.append('\n')
+          had_blank = True
+        line = sys.stdin.readline()
+        line_rstrip = line.rstrip()
+        if not (
+          len(line) and
+          line_rstrip[:len(indent)] == indent[:len(line_rstrip)] and
+          (len(line_rstrip) <= len(indent) or line[len(indent)] in '\t )')
+        ):
+          break
   else:
     match = re_element.match(line)
     if match is not None:
@@ -126,17 +144,12 @@ while len(line):
       begin = match.group(4)
 
       while len(stack) and stack[-1][0][:len(indent)] == indent:
-        #_, temp_class_name, _, _ = stack.pop()
-        #for temp_base_class, temp_fields in base_classes.pop().items():
-        #  base_classes[-1][
-        #    '{0:s}.{1:s}'.format(temp_class_name, temp_base_class)
-        #  ] = temp_fields
         temp_indent, _, _, _ = stack.pop()
         if len(delayed_line):
-          lines.append('{0:s};\n'.format(delayed_line[:-3]))
+          section2.append('{0:s};\n'.format(delayed_line[:-3]))
           delayed_line = ''
         else:
-          lines.append('{0:s}}};\n'.format(temp_indent))
+          section2.append('{0:s}}};\n'.format(temp_indent))
       _, class_name, base_class, full_base_class = stack[-1]
  
       fields = params.split(',')
@@ -144,14 +157,12 @@ while len(line):
         del fields[-1:]
       fields = [i.split() for i in fields]
       fields = [(type, name) for [type, name] in fields]
-      #i = len(base_classes[-2][class_name])
-      #base_classes[-2][class_name].extend(fields)
 
       if len(fields):
-        lines.append(delayed_line)
+        section2.append(delayed_line)
         delayed_line = ''
         for type, name in fields:
-          lines.append(
+          section2.append(
             '{0:s}{1:s} {2:s} = {3:s};\n'.format(
               indent,
               type,
@@ -172,13 +183,13 @@ while len(line):
     else:
       match = re_factory.match(line)
       if match is not None:
-        indent = '  ' #match.group(1)
+        indent = match.group(1)
         param = match.group(2)
         begin = match.group(3)
 
         if not had_blank:
-          sys.stdout.write('\n')
-        sys.stdout.write(
+          section3.append('\n')
+        section3.append(
           '''{0:s}def factory(tag, attrib = {{}}, *args, **kwargs):
 {1:s}  return tag_to_class.get(tag, {2:s})(tag, attrib, *args, **kwargs)
 '''.format(
@@ -247,51 +258,44 @@ while len(line):
               text = match.group(2)
  
               while len(stack) and stack[-1][0][:len(indent)] == indent:
-                #_, temp_class_name, _, _ = stack.pop()
-                #for temp_base_class, temp_fields in base_classes.pop().items():
-                #  base_classes[-1][
-                #    '{0:s}.{1:s}'.format(temp_class_name, temp_base_class)
-                #  ] = temp_fields
                 temp_indent, _, _, _ = stack.pop()
                 if len(delayed_line):
-                  lines.append('{0:s};\n'.format(delayed_line[:-3]))
+                  section2.append('{0:s};\n'.format(delayed_line[:-3]))
                   delayed_line = ''
                 else:
-                  lines.append('{0:s}}};\n'.format(temp_indent))
+                  section2.append('{0:s}}};\n'.format(temp_indent))
 
-              lines.append(delayed_line)
+              section2.append(delayed_line)
               delayed_line = ''
-              lines.append('{0:s}/* {1:s} */\n'.format(indent, text))
+              section2.append('{0:s}/* {1:s} */\n'.format(indent, text))
 
               line = sys.stdin.readline()
             else:
               line_rstrip = line.rstrip()
-              if len(line_rstrip):
+              if line_rstrip[:7] == 'import ' or line_rstrip[:8] == '#import ':
                 sys.stdout.write('  {0:s}\n'.format(line_rstrip))
+              elif len(line_rstrip):
+                section3.append(line_rstrip + '\n')
                 had_blank = False
               elif not had_blank:
-                sys.stdout.write('\n')
+                section3.append('\n')
                 had_blank = True
               line = sys.stdin.readline()
 
 while len(stack):
-  #_ , temp_class_name, _, _ = stack.pop()
-  #for temp_base_class, temp_fields in base_classes.pop().items():
-  #  base_classes[-1][
-  #    '{0:s}.{1:s}'.format(temp_class_name, temp_base_class)
-  #  ] = temp_fields
   temp_indent, _, _, _ = stack.pop()
   if len(delayed_line):
-    lines.append('{0:s};\n'.format(delayed_line[:-3]))
+    section2.append('{0:s};\n'.format(delayed_line[:-3]))
     delayed_line = ''
   else:
-    lines.append('{0:s}}};\n'.format(temp_indent))
-lines.append(delayed_line)
+    section2.append('{0:s}}};\n'.format(temp_indent))
+section2.append(delayed_line)
 
 sys.stdout.write(
-  '%}}\n\n%%\n\n{0:s}\n%%\n{1:s}'.format(
-    ''.join(lines),
+  '%}}\n\n%%\n\n{0:s}\n%%\n\n{1:s}'.format(
+    ''.join(section2),
     ''.join(
+      section3 +
       [
         '\n{0:s}{1:s}'.format(
           ''.join(lines),

diff --git a/expected.sed b/expected.sed
index 90b0f63..35b6412 100644 (file)
--- a/expected.sed
+++ b/expected.sed
@@ -1,3 +1,3 @@
-/^ *class/p
+/^ *class [A-Za-z_][A-Za-z0-9_]*(/p
 /^ *# GENERATE ELEMENT/,/^ *# GENERATE END/p
 /# GENERATE FACTORY/,/# GENERATE END/{s/^def factory/#&/; s/^  return tag_to_class/#&/; p}