git.polytechnique.org Git - xnet/blob - xnet/datasets/management/commands/loaddatasets.py

   1 # -*- coding: utf-8 -*-
   2
   3 """Load a specific dataset into the database."""
   4
   5 from django.core import management
   6 from django.db import transaction
   7 from django.utils import importlib
   8
   9
  10 class Command(management.BaseCommand):
  11     args = '<dataset_name dataset_name ...>'
  12     help = 'Load a specific dataset into the database.'
  13
  14     def handle(self, *dataset_names, **options):
  15         if not dataset_names:
  16             raise ValueError("At least one dataset name must be provided.")
  17         for name in dataset_names:
  18             self.load(name, **options)
  19
  20     def load(self, dataset_name, **options):
  21         """Loads a dataset."""
  22         verbosity = options.get('verbosity', 0)
  23         indent = options.pop('indent', '')
  24
  25         if verbosity > 0:
  26             self.stdout.write(u'%sImporting dataset %s...\n' % (indent, dataset_name))
  27
  28         dataset = importlib.import_module('xnet.datasets.%s' % dataset_name)
  29         setup = dataset.Setup(verbosity=verbosity)
  30
  31         # Handle required datasets
  32         for required in setup.requires:
  33             self.load(required, indent=indent + '  ', **options)
  34
  35         # Load the data
  36         with transaction.commit_on_success():
  37             setup.run()
  38
  39         if verbosity > 0:
  40             self.stdout.write(u'%sDone.\n' % indent)