Django snippets

''' Usage
    python rearrange_models.py inspected_models.py ordered_models.py
'''

import re
from cStringIO import StringIO

class DjangoModelsArrangement:
    def __init__(self):
        self.output = StringIO()
        self.classes = {}
        self.already_written = set()

    def get_classes(self, input):
        # This regular expression will divide the input as
        # [file_header, class_declaration1, class_body1, class_declaration2, class_body2, ...]
        parsed_classes = re.split(r'(class\s+.*?\(models\.Model\):)', input)
        self.header = parsed_classes[0]
        return [i+j for i, j in zip(parsed_classes[1::2], parsed_classes[2::2])]
    
    def parse_class(self, class_def):
        lines = class_def.split('\n')
        class_decl, class_body = lines[0], lines[1:]
        class_name = re.findall(r'class\s+(.*?)\(models\.Model\):', class_decl)[0]
        referenced_set = set()
        for line in lines:
            if 'ForeignKey' in line:
                referenced_class = re.findall(r'ForeignKey\((.*?)(?:,|\))', line)[0]
                if referenced_class and referenced_class != "'self'":
                    referenced_set.add(referenced_class)
        self.classes[class_name] = (referenced_set, class_def)
    
    def write_class(self, class_):
        if class_ not in self.already_written:
            self.already_written.add(class_)
            referenced_set, class_def = self.classes[class_]
            if referenced_set:
                for referenced_class in referenced_set:
                    self.write_class(referenced_class)
            self.output.write(class_def)
            
                
    def arrange(self, source_file, dest_file):
        input = open(source_file).read()
        classes = self.get_classes(input)
        for class_ in classes:
            self.parse_class(class_)
        for class_ in self.classes:
            self.write_class(class_)
        output_data = self.output.getvalue()
        open(dest_file, 'w').write(self.header + output_data)
        
if __name__ == '__main__':
    import sys
    if len(sys.argv) != 3:
        print __doc__
        sys.exit(1)
    DjangoModelsArrangement().arrange(sys.argv[1], sys.argv[2])