management.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. from django.apps import apps
  2. from django.db import DEFAULT_DB_ALIAS, router
  3. from django.db.models import signals
  4. from django.utils.encoding import smart_text
  5. from django.utils import six
  6. from django.utils.six.moves import input
  7. def update_contenttypes(app_config, verbosity=2, interactive=True, using=DEFAULT_DB_ALIAS, **kwargs):
  8. """
  9. Creates content types for models in the given app, removing any model
  10. entries that no longer have a matching model class.
  11. """
  12. if not app_config.models_module:
  13. return
  14. try:
  15. ContentType = apps.get_model('contenttypes', 'ContentType')
  16. except LookupError:
  17. return
  18. if not router.allow_migrate(using, ContentType):
  19. return
  20. ContentType.objects.clear_cache()
  21. app_label = app_config.label
  22. app_models = dict(
  23. (model._meta.model_name, model)
  24. for model in app_config.get_models())
  25. if not app_models:
  26. return
  27. # Get all the content types
  28. content_types = dict(
  29. (ct.model, ct)
  30. for ct in ContentType.objects.using(using).filter(app_label=app_label)
  31. )
  32. to_remove = [
  33. ct
  34. for (model_name, ct) in six.iteritems(content_types)
  35. if model_name not in app_models
  36. ]
  37. cts = [
  38. ContentType(
  39. name=smart_text(model._meta.verbose_name_raw),
  40. app_label=app_label,
  41. model=model_name,
  42. )
  43. for (model_name, model) in six.iteritems(app_models)
  44. if model_name not in content_types
  45. ]
  46. ContentType.objects.using(using).bulk_create(cts)
  47. if verbosity >= 2:
  48. for ct in cts:
  49. print("Adding content type '%s | %s'" % (ct.app_label, ct.model))
  50. # Confirm that the content type is stale before deletion.
  51. if to_remove:
  52. if kwargs.get('interactive', False):
  53. content_type_display = '\n'.join(
  54. ' %s | %s' % (ct.app_label, ct.model)
  55. for ct in to_remove
  56. )
  57. ok_to_delete = input("""The following content types are stale and need to be deleted:
  58. %s
  59. Any objects related to these content types by a foreign key will also
  60. be deleted. Are you sure you want to delete these content types?
  61. If you're unsure, answer 'no'.
  62. Type 'yes' to continue, or 'no' to cancel: """ % content_type_display)
  63. else:
  64. ok_to_delete = False
  65. if ok_to_delete == 'yes':
  66. for ct in to_remove:
  67. if verbosity >= 2:
  68. print("Deleting stale content type '%s | %s'" % (ct.app_label, ct.model))
  69. ct.delete()
  70. else:
  71. if verbosity >= 2:
  72. print("Stale content types remain.")
  73. def update_all_contenttypes(**kwargs):
  74. for app_config in apps.get_app_configs():
  75. update_contenttypes(app_config, **kwargs)
  76. signals.post_migrate.connect(update_contenttypes)
  77. if __name__ == "__main__":
  78. update_all_contenttypes()