def _cleanup():
"""Reset indexes, clean all connections cached in the _document_registry.
Drop databases
"""
from mongoengine.queryset import QuerySet
from mongoengine.base import _document_registry
QuerySet._reset_already_indexed()
for k, doc in _document_registry.items():
if not hasattr(doc, '_get_collection'):
continue
if doc._meta.get('abstract', False):
continue
col = doc._get_collection()
col.database.connection.drop_database(col.database.name)
def import_objects(options, style):
from django.apps import apps
from django import setup
if not apps.ready:
setup()
from django.conf import settings
dont_load_cli = options.get('dont_load', [])
dont_load_conf = getattr(settings, 'SHELL_PLUS_DONT_LOAD', [])
dont_load = dont_load_cli + dont_load_conf
dont_load_any_models = '*' in dont_load
quiet_load = options.get('quiet_load')
model_aliases = getattr(settings, 'SHELL_PLUS_MODEL_ALIASES', {})
app_prefixes = getattr(settings, 'SHELL_PLUS_APP_PREFIXES', {})
SHELL_PLUS_PRE_IMPORTS = getattr(settings, 'SHELL_PLUS_PRE_IMPORTS', {})
imported_objects = {}
load_models = {}
def get_dict_from_names_to_possible_models(): # type: () -> Dict[str, List[str]]
"""
Collect dictionary from names to possible models. Model is represented as his full path.
Name of model can be alias if SHELL_PLUS_MODEL_ALIASES or SHELL_PLUS_APP_PREFIXES is specified for this model.
This dictionary is used by collision resolver.
At this phase we can't import any models, because collision resolver can change results.
:return: Dict[str, List[str]]. Key is name, value is list of full model's path's.
"""
models_to_import = {} # type: Dict[str, List[str]]
for app_mod, models in sorted(load_models.items()):
app_name = get_app_name(app_mod)
app_aliases = model_aliases.get(app_name, {})
prefix = app_prefixes.get(app_name)
for model_name in sorted(models):
if "%s.%s" % (app_name, model_name) in dont_load:
continue
alias = app_aliases.get(model_name)
if not alias:
if prefix:
alias = "%s_%s" % (prefix, model_name)
else:
alias = model_name
models_to_import.setdefault(alias, [])
models_to_import[alias].append("%s.%s" % (app_mod, model_name))
return models_to_import
def import_subclasses():
base_classes_to_import = getattr(settings, 'SHELL_PLUS_SUBCLASSES_IMPORT', []) # type: List[Union[str, type]]
if base_classes_to_import:
if not quiet_load:
print(style.SQL_TABLE("# Shell Plus Subclasses Imports"))
perform_automatic_imports(SubclassesFinder(base_classes_to_import).collect_subclasses())
def import_models():
"""
Perform collision resolving and imports all models.
When collisions are resolved we can perform imports and print information's, because it is last phase.
This function updates imported_objects dictionary.
"""
modules_to_models = CollisionResolvingRunner().run_collision_resolver(get_dict_from_names_to_possible_models())
perform_automatic_imports(modules_to_models)
def perform_automatic_imports(modules_to_classes): # type: (Dict[str, List[Tuple[str, str]]]) -> ()
"""
Import elements from given dictionary.
:param modules_to_classes: dictionary from module name to tuple.
First element of tuple is model name, second is model alias.
If both elements are equal than element is imported without alias.
"""
for full_module_path, models in modules_to_classes.items():
model_labels = []
for (model_name, alias) in sorted(models):
try:
imported_objects[alias] = import_string("%s.%s" % (full_module_path, model_name))
if model_name == alias:
model_labels.append(model_name)
else:
model_labels.append("%s (as %s)" % (model_name, alias))
except ImportError as e:
if options.get("traceback"):
traceback.print_exc()
if not options.get('quiet_load'):
print(style.ERROR(
"Failed to import '%s' from '%s' reason: %s" % (model_name, full_module_path, str(e))))
if not options.get('quiet_load'):
print(style.SQL_COLTYPE("from %s import %s" % (full_module_path, ", ".join(model_labels))))
def get_apps_and_models():
for app in apps.get_app_configs():
if app.models_module:
yield app.models_module, app.get_models()
mongoengine = False
try:
from mongoengine.base import _document_registry
mongoengine = True
except ImportError:
pass
# Perform pre-imports before any other imports
if SHELL_PLUS_PRE_IMPORTS:
if not quiet_load:
print(style.SQL_TABLE("# Shell Plus User Pre Imports"))
imports = import_items(SHELL_PLUS_PRE_IMPORTS, style, quiet_load=quiet_load)
for k, v in imports.items():
imported_objects[k] = v
if mongoengine and not dont_load_any_models:
for name, mod in _document_registry.items():
name = name.split('.')[-1]
app_name = get_app_name(mod.__module__)
if app_name in dont_load or ("%s.%s" % (app_name, name)) in dont_load:
continue
load_models.setdefault(mod.__module__, [])
load_models[mod.__module__].append(name)
if not dont_load_any_models:
for app_mod, app_models in get_apps_and_models():
if not app_models:
continue
app_name = get_app_name(app_mod.__name__)
if app_name in dont_load:
continue
for mod in app_models:
if "%s.%s" % (app_name, mod.__name__) in dont_load:
continue
if mod.__module__:
# Only add the module to the dict if `__module__` is not empty.
load_models.setdefault(mod.__module__, [])
load_models[mod.__module__].append(mod.__name__)
import_subclasses()
if not quiet_load:
print(style.SQL_TABLE("# Shell Plus Model Imports%s") % (' SKIPPED' if dont_load_any_models else ''))
import_models()
# Imports often used from Django
if getattr(settings, 'SHELL_PLUS_DJANGO_IMPORTS', True):
if not quiet_load:
print(style.SQL_TABLE("# Shell Plus Django Imports"))
imports = import_items(SHELL_PLUS_DJANGO_IMPORTS, style, quiet_load=quiet_load)
for k, v in imports.items():
imported_objects[k] = v
SHELL_PLUS_IMPORTS = getattr(settings, 'SHELL_PLUS_IMPORTS', {})
if SHELL_PLUS_IMPORTS:
if not quiet_load:
print(style.SQL_TABLE("# Shell Plus User Imports"))
imports = import_items(SHELL_PLUS_IMPORTS, style, quiet_load=quiet_load)
for k, v in imports.items():
imported_objects[k] = v
# Perform post-imports after any other imports
SHELL_PLUS_POST_IMPORTS = getattr(settings, 'SHELL_PLUS_POST_IMPORTS', {})
if SHELL_PLUS_POST_IMPORTS:
if not quiet_load:
print(style.SQL_TABLE("# Shell Plus User Post Imports"))
imports = import_items(SHELL_PLUS_POST_IMPORTS, style, quiet_load=quiet_load)
for k, v in imports.items():
imported_objects[k] = v
return imported_objects
def import_objects(options, style):
# XXX: (Temporary) workaround for ticket #1796: force early loading of all
# models from installed apps. (this is fixed by now, but leaving it here
# for people using 0.96 or older trunk (pre [5919]) versions.
from django.db.models.loading import get_models, get_apps
mongoengine = False
try:
from mongoengine.base import _document_registry
mongoengine = True
except:
pass
loaded_models = get_models() # NOQA
from django.conf import settings
imported_objects = {'settings': settings}
dont_load_cli = options.get('dont_load') # optparse will set this to [] if it doensnt exists
dont_load_conf = getattr(settings, 'SHELL_PLUS_DONT_LOAD', [])
dont_load = dont_load_cli + dont_load_conf
quiet_load = options.get('quiet_load')
model_aliases = getattr(settings, 'SHELL_PLUS_MODEL_ALIASES', {})
# Perform pre-imports before any other imports
imports = import_items(getattr(settings, 'SHELL_PLUS_PRE_IMPORTS', {}))
for k, v in imports.items():
imported_objects[k] = v
load_models = {}
if mongoengine:
for name, mod in _document_registry.items():
name = name.split('.')[-1]
app_name = mod.__module__.split('.')[-2]
if app_name in dont_load or ("%s.%s" % (app_name, name)) in dont_load:
continue
load_models.setdefault(mod.__module__, [])
load_models[mod.__module__].append(name)
for app_mod in get_apps():
app_models = get_models(app_mod)
if not app_models:
continue
app_name = app_mod.__name__.split('.')[-2]
if app_name in dont_load:
continue
app_aliases = model_aliases.get(app_name, {})
for mod in app_models:
if "%s.%s" % (app_name, mod.__name__) in dont_load:
continue
load_models.setdefault(mod.__module__, [])
load_models[mod.__module__].append(mod.__name__)
for app_mod, models in load_models.items():
app_name = app_mod.split('.')[-2]
app_aliases = model_aliases.get(app_name, {})
model_labels = []
for model_name in models:
try:
imported_object = getattr(__import__(app_mod, {}, {}, model_name), model_name)
if "%s.%s" % (app_name, model_name) in dont_load:
continue
alias = app_aliases.get(model_name, model_name)
imported_objects[alias] = imported_object
if model_name == alias:
model_labels.append(model_name)
else:
model_labels.append("%s (as %s)" % (model_name, alias))
except AttributeError as e:
if options.get("traceback"):
traceback.print_exc()
if not quiet_load:
print(style.ERROR("Failed to import '%s' from '%s' reason: %s" % (model_name, app_mod, str(e))))
continue
if not quiet_load:
print(style.SQL_COLTYPE("From '%s' autoload: %s" % (app_mod.split('.')[-2], ", ".join(model_labels))))
# Perform post-imports after any other imports
imports = import_items(getattr(settings, 'SHELL_PLUS_POST_IMPORTS', {}))
for k, v in imports.items():
imported_objects[k] = v
return imported_objects
