Fixed adding regular tags.
[zzz-floof.git] / 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']