0002_fix_space_foreign_keys.py 文件源码

def fix_space_foreign_keys(apps, schema_editor):
    models = ('Area', 'Stair', 'LineObstacle', 'Obstacle')
    Space = apps.get_model('mapdata', 'Space')
    MAP_PATH = os.path.join(settings.DATA_DIR, 'map')
    PACKAGE_PATHS = [os.path.join(MAP_PATH, dirname) for dirname in os.listdir(MAP_PATH)]
    PACKAGE_PATHS = [path for path in PACKAGE_PATHS if os.path.isdir(path)]
    spaces = {}
    for space in Space.objects.all():
        spaces.setdefault(space.section.name + '_' + space.level, []).append(space)
    for model_name in models:
        model = apps.get_model('mapdata', model_name)
        orig_dir_names = [model._meta.default_related_name.lower()]
        if model.__name__ == 'Area':
            orig_dir_names = ['arealocations', 'stuffedareas']

        orig_objects = []
        for package_path in PACKAGE_PATHS:
            for orig_dir_name in orig_dir_names:
                dir_name = os.path.join(package_path, orig_dir_name)
                if not os.path.isdir(dir_name):
                    continue
                for filename in os.listdir(dir_name):
                    abs_filename = os.path.join(dir_name, filename)
                    if not filename.endswith('.json') or not os.path.isfile(abs_filename):
                        continue
                    obj = json.load(open(abs_filename))
                    obj['name'] = filename[:-5]
                    obj['geometry'] = shape(obj['geometry'])
                    orig_objects.append(obj)
        matches = {}
        for obj in model.objects.all().order_by('id' if hasattr(model, 'id') else 'locationslug_ptr_id'):
            for i, orig_obj in enumerate(orig_objects):
                if obj.geometry.almost_equals(orig_obj['geometry']):
                    matches.setdefault(i, []).append(obj)
                    break

        for orig_i, objects in matches.items():
            orig_obj = orig_objects[orig_i]

            if '-' in orig_obj['level']:
                splitted = orig_obj['level'].split('-')
                possible_spaces = spaces[splitted[0] + '_upper'] + spaces[splitted[1] + '_lower']
            else:
                possible_spaces = spaces[orig_obj['level'] + '_']

            possible_spaces = [space for space in possible_spaces
                               if space.geometry.intersects(orig_obj['geometry'])]
            if len(objects) == len(possible_spaces):
                pass  # nice
            elif len(objects) > len(possible_spaces):
                pass  # well, whatever
            else:
                def compare(space):
                    return space.geometry.intersection(orig_obj['geometry']).area

                possible_spaces.sort(key=compare, reverse=True)

            for i, obj in enumerate(objects[:len(possible_spaces)]):
                obj.space = possible_spaces[i]
                obj.save()