floof/lib/fixtures.py

   1 import types
   2 import sys
   3 import os
   4 import simplejson
   5
   6 from floof import model as model
   7
   8 """
   9 This module can be used for loading data into your models, for example when setting up default application data,
  10 unit tests, JSON export/import and importing/exporting legacy data. Data is serialized to and from the JSON format.
  11 """
  12
  13 VALID_FIXTURE_FILE_EXTENSIONS = ['.json']
  14
  15 def load_data(model, filename=None, base_dir=None):
  16     """Installs provided fixture files into given model. Filename may be directory, file or list of dirs or files. If filename is
  17     None, assumes that source file is located in fixtures/model_module_name/model_tablename.yaml of your application directory,
  18     for example MyProject/fixtures/news/newsitems.yaml. The base_dir argument is the top package of the application unless
  19     specified. You can also pass the name of a table instead of a model class."""
  20
  21     if type(model) is types.StringType:
  22         return load_data_to_table(model, filename, base_dir)
  23     else:
  24         if filename is None:
  25             filename = _default_fixture_path_for_model(model, base_dir)
  26         return _load_data_from_file(model, filename)
  27
  28 def load_data_to_table(table, filename=None, base_dir=None):
  29     """Installs data directly into a table. Useful if table does not have a corresponding model, for example a many-to-many join table.
  30     """
  31
  32     if filename is None:
  33         filename = _default_fixture_path_for_table(table, base_dir)
  34     _load_data_to_table(table, filename)
  35
  36 def dump_data(model, filename=None, **params):
  37     """Dumps data to given destination. Params are optional arguments for selecting data. If filename is None, assumes that destination
  38     file is located in fixtures/model_module_name/model_name_lowercase.yaml of your application directory, for example
  39     MyProject/fixtures/news/newsitem.yaml.
  40     """
  41
  42     if filename is None:
  43         filename = _default_fixture_path_for_model(model)
  44     _dump_data_to_file(model, filename, **params)
  45
  46 _base_dir = os.path.dirname(os.path.dirname(__file__))
  47
  48 def _default_fixture_path_for_model(model, base_dir=None):
  49     if base_dir is None:
  50         base_dir = _base_dir
  51     path = os.path.join(base_dir, 'fixtures')
  52     module_dirs = model.__module__.split('.', 2)[-1].split('.')
  53     for dir in module_dirs:
  54         path = os.path.join(path, dir)
  55     return os.path.join(path, model.table.name + '.json')
  56
  57 def _default_fixture_path_for_table(table, base_dir=None):
  58     if base_dir is None:
  59         base_dir = _base_dir
  60     module_dirs = table.split('.')
  61     path = os.path.join(base_dir, 'fixtures')
  62     for name in module_dirs:
  63         path = os.path.join(path, name)
  64     return path + ".json"
  65
  66 def _is_fixture_file(filename):
  67     basename, ext = os.path.splitext(filename)
  68     return (ext.lower() in VALID_FIXTURE_FILE_EXTENSIONS)
  69
  70 def _load_data_from_dir(model, dirname):
  71     for dirpath, dirnames, filenames in os.walk(dirname):
  72         for filename in filenames:
  73             _load_data_from_file(model, filename)
  74
  75 def _load_data_from_file(model, filename):
  76     if not _is_fixture_file(filename):
  77         return
  78     fp = file(filename, 'r')
  79     data = simplejson.load(fp)
  80     fp.close()
  81     retval = None
  82     if type(data) is types.ListType:
  83         retval = []
  84         for item in data:
  85             retval.append(_load_instance_from_dict(model, item))
  86     elif type(data) is types.DictType:
  87         retval = {}
  88         for key, item in data.iteritems():
  89             retval[key] = _load_instance_from_dict(model, item)
  90     return retval
  91
  92 def _load_data_to_table(tablename, filename):
  93     if not _is_fixture_file(filename):
  94         return
  95     fp = file(filename, 'r')
  96     data = simplejson.load(fp)
  97     fp.close()
  98     tablename = tablename.split(".")[-1]
  99     table = model.context.metadata.tables[tablename]
 100     if type(data) is types.ListType:
 101         for item in data:
 102             table.insert(item).execute()
 103     elif type(data) is types.DictType:
 104         for key, item in data.iteritems():
 105             table.insert(item).execute()
 106     return data
 107
 108 def _dump_data_to_file(model, filename, **params):
 109     if params:
 110         queryset = model.select_by(**params)
 111     else:
 112         queryset = model.select()
 113     data = []
 114     for instance in queryset:
 115         data.append(_dump_instance_to_dict(instance))
 116     fp = file(filename, 'w')
 117     simplejson.dump(data, fp)
 118     fp.close()
 119
 120 def _load_instance_from_dict(model, dict):
 121     if not dict: return
 122     instance = model()
 123     fields = model._descriptor.fields.keys()
 124     for k, v in dict.iteritems():
 125         if k in fields:
 126             setattr(instance, k, v)
 127     instance.flush()
 128     return instance
 129
 130 def _dump_instance_to_dict(instance):
 131     if hasattr(instance, 'to_json'):
 132         return instance.to_json()
 133     d = {}
 134     fields = instance._descriptor.fields.keys()
 135     for field in fields:
 136         d[field] = getattr(instance, field)
 137     return d
 138
 139 __all__ = ['load_data', 'dump_data']