single_set_manager.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. from .manager import manager
  2. from .single_set_model import single_set_model
  3. class single_set_manager(manager):
  4. def __init_subclass__(
  5. cls,
  6. target_model: type,
  7. related_models: str | tuple | list
  8. ) -> None:
  9. super().__init_subclass__()
  10. if not issubclass(target_model, single_set_model):
  11. raise TypeError("Target must be subclass of single_set_model.")
  12. if type(related_models) is str:
  13. related_models = tuple([ related_models ])
  14. if type(related_models) is list:
  15. related_models = tuple(related_models)
  16. if type(related_models) is not tuple:
  17. raise TypeError("Related model must be str, list or tuple.")
  18. cls.__related_models = related_models
  19. cls.__target_model = target_model
  20. @classmethod
  21. def _get_related_models(cls) -> tuple:
  22. return cls.__related_models
  23. @classmethod
  24. def _get_model(cls) -> type:
  25. return cls.__target_model
  26. @classmethod
  27. async def add(cls, target: str) -> single_set_model:
  28. target_model = cls._get_model()
  29. result = await target_model \
  30. .filter(content = target) \
  31. .first()
  32. if result is None:
  33. return await cls._create(target)
  34. return result
  35. @classmethod
  36. async def _create(cls, content: str) -> single_set_model:
  37. result = cls \
  38. ._get_model() \
  39. .get_proxy() \
  40. .create(content) \
  41. .result()
  42. await result.save()
  43. return result
  44. @classmethod
  45. async def all(cls) -> tuple:
  46. return tuple(await cls._get_model().all())
  47. @classmethod
  48. async def clean(cls, target: str) -> None:
  49. target_model = cls._get_model()
  50. related_models = cls._get_related_models()
  51. target = await target_model \
  52. .filter(content = target) \
  53. .first()
  54. await target.fetch_related(*related_models)
  55. for count in related_models:
  56. if len(target.__getattribute__(count)) > 1:
  57. return
  58. await target.delete()