Add datasets module.

[xnet] / xnet / datasets / management / commands / loaddatasets.py

diff --git a/xnet/datasets/management/commands/loaddatasets.py b/xnet/datasets/management/commands/loaddatasets.py
new file mode 100644 (file)
index 0000000..75c3b66
--- /dev/null
+++ b/xnet/datasets/management/commands/loaddatasets.py
@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+
+"""Load a specific dataset into the database."""
+
+from django.core import management
+from django.db import transaction
+from django.utils import importlib
+
+
+class Command(management.BaseCommand):
+    args = '<dataset_name dataset_name ...>'
+    help = 'Load a specific dataset into the database.'
+
+    def handle(self, *dataset_names, **options):
+        if not dataset_names:
+            raise ValueError("At least one dataset name must be provided.")
+        for name in dataset_names:
+            self.load(name, **options)
+
+    def load(self, dataset_name, **options):
+        """Loads a dataset."""
+        verbosity = options.get('verbosity', 0)
+        indent = options.pop('indent', '')
+
+        if verbosity > 0:
+            self.stdout.write(u'%sImporting dataset %s...\n' % (indent, dataset_name))
+
+        dataset = importlib.import_module('xnet.datasets.%s' % dataset_name)
+        setup = dataset.Setup(verbosity=verbosity)
+
+        # Handle required datasets
+        for required in setup.requires:
+            self.load(required, indent=indent + '  ', **options)
+
+        # Load the data
+        with transaction.commit_on_success():
+            setup.run()
+
+        if verbosity > 0:
+            self.stdout.write(u'%sDone.\n' % indent)