[SQL-CVS] r1234 - SQLObject/trunk/sqlobject/util
SQLObject is a Python ORM.
Brought to you by:
ianbicking,
phd
|
From: <sub...@co...> - 2005-11-09 22:20:44
|
Author: test
Date: 2005-11-09 22:20:39 +0000 (Wed, 09 Nov 2005)
New Revision: 1234
Added:
SQLObject/trunk/sqlobject/util/csvimport.py
Log:
Added a utility to import data from a CSV file, primarily intended for test fixtures
Added: SQLObject/trunk/sqlobject/util/csvimport.py
===================================================================
--- SQLObject/trunk/sqlobject/util/csvimport.py 2005-11-09 22:20:06 UTC (rev 1233)
+++ SQLObject/trunk/sqlobject/util/csvimport.py 2005-11-09 22:20:39 UTC (rev 1234)
@@ -0,0 +1,349 @@
+"""
+Import from a CSV file or directory of files.
+
+CSV files should have a header line that lists columns. Headers can
+also be appended with ``:type`` to indicate the type of the field.
+``escaped`` is the default, though it can be overridden by the importer.
+Supported types:
+
+``:python``:
+ A python expression, run through ``eval()``. This can be a
+ security risk, pass in ``allow_python=False`` if you don't want to
+ allow it.
+
+``:int``:
+ Integer
+
+``:float``:
+ Float
+
+``:str``:
+ String
+
+``:escaped``:
+ A string with backslash escapes (note that you don't put quotation
+ marks around the value)
+
+``:base64``:
+ A base64-encoded string
+
+``:date``:
+ ISO date, like YYYY-MM-DD; this can also be ``NOW+days`` or
+ ``NOW-days``
+
+``:datetime``:
+ ISO date/time like YYYY-MM-DDTHH:MM:SS (either T or a space can be
+ used to separate the time, and seconds are optional). This can
+ also be ``NOW+seconds`` or ``NOW-seconds``
+
+``:bool``:
+ Converts true/false/yes/no/on/off/1/0 to boolean value
+
+``:ref``:
+ This will be resolved to the ID of the object named in this column
+ (None if the column is empty). @@: Since there's no ordering,
+ there's no way to promise the object already exists.
+
+You can also get back references to the objects if you have a special
+``[name]`` column.
+
+Any column named ``[comment]`` or with no name will be ignored.
+
+In any column you can put ``[default]`` to exclude the value and use
+whatever default the class wants. ``[null]`` will use NULL.
+
+Lines that begin with ``[comment]`` are ignored.
+"""
+
+from datetime import datetime, date, timedelta
+import os
+import csv
+import types
+
+__all__ = ['load_csv_from_directory',
+ 'load_csv',
+ 'create_data']
+
+DEFAULT_TYPE = 'escaped'
+
+def create_data(data, class_getter, keyorder=None):
+ """
+ Create the ``data``, which is the return value from
+ ``load_csv()``. Classes will be resolved with the callable
+ ``class_getter``; or if ``class_getter`` is a module then the
+ class names will be attributes of that.
+
+ Returns a dictionary of ``{object_name: object(s)}``, using the
+ names from the ``[name]`` columns (if there are any). If a name
+ is used multiple times, you get a list of objects, not a single
+ object.
+
+ If ``keyorder`` is given, then the keys will be retrieved in that
+ order. It can be a list/tuple of names, or a sorting function.
+ If not given and ``class_getter`` is a module and has a
+ ``soClasses`` function, then that will be used for the order.
+ """
+ objects = {}
+ classnames = data.keys()
+ if (not keyorder and isinstance(class_getter, types.ModuleType)
+ and hasattr(class_getter, 'soClasses')):
+ keyorder = [c.__name__ for c in class_getter.soClasses]
+ if not keyorder:
+ classnames.sort()
+ elif isinstance(keyorder, (list, tuple)):
+ all = classnames
+ classnames = [name for name in keyorder if name in classnames]
+ for name in all:
+ if name not in classnames:
+ classnames.append(name)
+ else:
+ classnames.sort(keyorder)
+ for classname in classnames:
+ items = data[classname]
+ if not items:
+ continue
+ if isinstance(class_getter, types.ModuleType):
+ soClass = getattr(class_getter, classname)
+ else:
+ soClass = class_getter(classname)
+ for item in items:
+ for key, value in item.items():
+ if isinstance(value, Reference):
+ resolved = objects.get(value.name)
+ if not resolved:
+ raise ValueError(
+ "Object reference to %r does not have target"
+ % value.name)
+ elif (isinstance(resolved, list)
+ and len(resolved) > 1):
+ raise ValueError(
+ "Object reference to %r is ambiguous (got %r)"
+ % (value.name, resolved))
+ item[key] = resolved.id
+ if '[name]' in item:
+ name = item.pop('[name]').strip()
+ else:
+ name = None
+ inst = soClass(**item)
+ if name:
+ if name in objects:
+ if isinstance(objects[name], list):
+ objects[name].append(inst)
+ else:
+ objects[name] = [objects[name], inst]
+ else:
+ objects[name] = inst
+ return objects
+
+
+def load_csv_from_directory(directory,
+ allow_python=True, default_type=DEFAULT_TYPE,
+ allow_multiple_classes=True):
+ """
+ Load the data from all the files in a directory. Filenames
+ indicate the class, with ``general.csv`` for data not associated
+ with a class. Return data just like ``load_csv`` does.
+
+ This might cause problems on case-insensitive filesystems.
+ """
+ results = {}
+ for filename in os.listdir(directory):
+ base, ext = os.path.splitext(filename)
+ if ext.lower() != '.csv':
+ continue
+ f = open(os.path.join(directory, filename), 'rb')
+ csvreader = csv.reader(f)
+ data = load_csv(csvreader, allow_python=allow_python,
+ default_type=default_type,
+ default_class=base,
+ allow_multiple_classes=allow_multiple_classes)
+ f.close()
+ for classname, items in data.items():
+ results.setdefault(classname, []).extend(items)
+ return results
+
+def load_csv(csvreader, allow_python=True, default_type=DEFAULT_TYPE,
+ default_class=None, allow_multiple_classes=True):
+ """
+ Loads the CSV file, returning a list of dictionaries with types
+ coerced.
+ """
+ current_class = default_class
+ current_headers = None
+ results = {}
+
+ for row in csvreader:
+ if not [cell for cell in row if cell.strip()]:
+ # empty row
+ continue
+
+ if row and row[0].strip() == 'CLASS:':
+ if not allow_multiple_classes:
+ raise ValueError(
+ "CLASS: line in CSV file, but multiple classes are not allowed in this file (line: %r)"
+ % row)
+ if not row[1:]:
+ raise ValueError(
+ "CLASS: in line in CSV file, with no class name in next column (line: %r)"
+ % row)
+ current_class = row[1]
+ current_headers = None
+ continue
+
+ if not current_class:
+ raise ValueError(
+ "No CLASS: line given, and there is no default class for this file (line: %r"
+ % row)
+
+ if current_headers is None:
+ current_headers = _parse_headers(row, default_type)
+ continue
+
+ if row[0] == '[comment]':
+ continue
+
+ # Pad row with empty strings:
+ row += ['']*(len(current_headers) - len(row))
+ row_converted = {}
+ for value, (name, coercer, args) in zip(row, current_headers):
+ if name is None:
+ # Comment
+ continue
+ if value == '[default]':
+ continue
+ if value == '[null]':
+ row_converted[name] = None
+ continue
+ args = (value,) + args
+ row_converted[name] = coercer(*args)
+
+ results.setdefault(current_class, []).append(row_converted)
+
+ return results
+
+def _parse_headers(header_row, default_type):
+ headers = []
+ for name in header_row:
+ original_name = name
+ if ':' in name:
+ name, type = name.split(':', 1)
+ else:
+ type = default_type
+ if type == 'python' and not allow_python:
+ raise ValueError(
+ ":python header given when python headers are not allowed (with header %r"
+ % original_name)
+ name = name.strip()
+ if name == '[comment]' or not name:
+ headers.append((None, None, None))
+ continue
+ type = type.strip().lower()
+ if '(' in type:
+ type, arg = type.split('(', 1)
+ if not arg.endswith(')'):
+ raise ValueError(
+ "Arguments (in ()'s) do not end with ): %r"
+ % original_name)
+ args = (arg[:-1],)
+ else:
+ args = ()
+ if name == '[name]':
+ type = 'str'
+ coercer, args = get_coercer(type)
+ headers.append((name, coercer, args))
+ return headers
+
+_coercers = {}
+def get_coercer(type):
+ if type not in _coercers:
+ raise ValueError(
+ "Coercion type %r not known (I know: %s)"
+ % (type, ', '.join(_coercers.keys())))
+ return _coercers[type]
+
+def register_coercer(type, coercer, *args):
+ _coercers[type] = (coercer, args)
+
+def identity(v):
+ return v
+
+register_coercer('str', identity)
+register_coercer('string', identity)
+
+def decode_string(v, encoding):
+ return v.decode(encoding)
+
+register_coercer('escaped', decode_string, 'string_escape')
+register_coercer('strescaped', decode_string, 'string_escape')
+register_coercer('base64', decode_string, 'base64')
+
+register_coercer('int', int)
+register_coercer('float', float)
+
+def parse_python(v):
+ return eval(v, {}, {})
+
+register_coercer('python', parse_python)
+
+def parse_date(v):
+ v = v.strip()
+ if not v:
+ return None
+ if v.startswith('NOW-') or v.startswith('NOW+'):
+ days = int(v[3:])
+ now = date.today()
+ return now+timedelta(days)
+ else:
+ parsed = time.strptime(v, '%Y-%m-%d')
+ return date.fromtimestamp(time.mktime(parsed))
+
+register_coercer('date', parse_date)
+
+def parse_datetime(v):
+ v = v.strip()
+ if not v:
+ return None
+ if v.startswith('NOW-') or v.startswith('NOW+'):
+ seconds = int(v[3:])
+ now = datetime.now()
+ return now+timedelta(0, seconds)
+ else:
+ fmts = ['%Y-%m-%dT%H:%M:%S',
+ '%Y-%m-%d %H:%M:%S',
+ '%Y-%m-%dT%H:%M',
+ '%Y-%m-%d %H:%M']
+ for fmt in fmts[:-1]:
+ try:
+ parsed = time.strptime(v, fmt)
+ break
+ except ValueError:
+ pass
+ else:
+ parsed = time.strptime(v, fmts[-1])
+ return datetime.fromtimestamp(time.mktime(parsed))
+
+register_coercer('datetime', parse_datetime)
+
+class Reference(object):
+ def __init__(self, name):
+ self.name = name
+
+def parse_ref(v):
+ if not v.strip():
+ return None
+ else:
+ return Reference(v)
+
+register_coercer('ref', parse_ref)
+
+def parse_bool(v):
+ v = v.strip().lower()
+ if v in ('y', 'yes', 't', 'true', 'on', '1'):
+ return True
+ elif v in ('n', 'no', 'f', 'false', 'off', '0'):
+ return False
+ raise ValueError(
+ "Value is not boolean-like: %r" % value)
+
+register_coercer('bool', parse_bool)
+register_coercer('boolean', parse_bool)
|