sorting.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. import re
  2. from collections.abc import Callable, Iterable
  3. from typing import TYPE_CHECKING, Any
  4. if TYPE_CHECKING:
  5. from .settings import Config
  6. else:
  7. Config = Any
  8. _import_line_intro_re = re.compile("^(?:from|import) ")
  9. _import_line_midline_import_re = re.compile(" import ")
  10. def module_key(
  11. module_name: str,
  12. config: Config,
  13. sub_imports: bool = False,
  14. ignore_case: bool = False,
  15. section_name: Any | None = None,
  16. straight_import: bool | None = False,
  17. ) -> str:
  18. match = re.match(r"^(\.+)\s*(.*)", module_name)
  19. if match:
  20. sep = " " if config.reverse_relative else "_"
  21. module_name = sep.join(match.groups())
  22. prefix = ""
  23. if ignore_case:
  24. module_name = str(module_name).lower()
  25. else:
  26. module_name = str(module_name)
  27. if sub_imports and config.order_by_type:
  28. if module_name in config.constants:
  29. prefix = "A"
  30. elif module_name in config.classes:
  31. prefix = "B"
  32. elif module_name in config.variables:
  33. prefix = "C"
  34. elif module_name.isupper() and len(module_name) > 1: # see issue #376
  35. prefix = "A"
  36. elif module_name in config.classes or module_name[0:1].isupper():
  37. prefix = "B"
  38. else:
  39. prefix = "C"
  40. if not config.case_sensitive:
  41. module_name = module_name.lower()
  42. length_sort = (
  43. config.length_sort
  44. or (config.length_sort_straight and straight_import)
  45. or str(section_name).lower() in config.length_sort_sections
  46. )
  47. _length_sort_maybe = (str(len(module_name)) + ":" + module_name) if length_sort else module_name
  48. return f"{(module_name in config.force_to_top and 'A') or 'B'}{prefix}{_length_sort_maybe}"
  49. def section_key(line: str, config: Config) -> str:
  50. section = "B"
  51. if (
  52. not config.sort_relative_in_force_sorted_sections
  53. and config.reverse_relative
  54. and line.startswith("from .")
  55. ):
  56. match = re.match(r"^from (\.+)\s*(.*)", line)
  57. if match: # pragma: no cover - regex always matches if line starts with "from ."
  58. line = f"from {' '.join(match.groups())}"
  59. if config.group_by_package and line.strip().startswith("from"):
  60. line = line.split(" import ", 1)[0]
  61. if config.lexicographical:
  62. line = _import_line_intro_re.sub("", _import_line_midline_import_re.sub(".", line))
  63. else:
  64. line = re.sub("^from ", "", line)
  65. line = re.sub("^import ", "", line)
  66. if config.sort_relative_in_force_sorted_sections:
  67. sep = " " if config.reverse_relative else "_"
  68. line = re.sub(r"^(\.+)", rf"\1{sep}", line)
  69. if line.split(" ")[0] in config.force_to_top:
  70. section = "A"
  71. # * If honor_case_in_force_sorted_sections is true, and case_sensitive and
  72. # order_by_type are different, only ignore case in part of the line.
  73. # * Otherwise, let order_by_type decide the sorting of the whole line. This
  74. # is only "correct" if case_sensitive and order_by_type have the same value.
  75. if config.honor_case_in_force_sorted_sections and config.case_sensitive != config.order_by_type:
  76. split_module = line.split(" import ", 1)
  77. if len(split_module) > 1:
  78. module_name, names = split_module
  79. if not config.case_sensitive:
  80. module_name = module_name.lower()
  81. if not config.order_by_type:
  82. names = names.lower()
  83. line = f"{module_name} import {names}"
  84. elif not config.case_sensitive:
  85. line = line.lower()
  86. elif not config.order_by_type:
  87. line = line.lower()
  88. return f"{section}{len(line) if config.length_sort else ''}{line}"
  89. def sort(
  90. config: Config,
  91. to_sort: Iterable[str],
  92. key: Callable[[str], Any] | None = None,
  93. reverse: bool = False,
  94. ) -> list[str]:
  95. return config.sorting_function(to_sort, key=key, reverse=reverse)
  96. def naturally(
  97. to_sort: Iterable[str], key: Callable[[str], Any] | None = None, reverse: bool = False
  98. ) -> list[str]:
  99. """Returns a naturally sorted list"""
  100. if key is None:
  101. key_callback = _natural_keys
  102. else:
  103. def key_callback(text: str) -> list[Any]:
  104. return _natural_keys(key(text))
  105. return sorted(to_sort, key=key_callback, reverse=reverse)
  106. def _atoi(text: str) -> Any:
  107. return int(text) if text.isdigit() else text
  108. def _natural_keys(text: str) -> list[Any]:
  109. return [_atoi(c) for c in re.split(r"(\d+)", text)]