123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113 |
- import ast
- from pprint import PrettyPrinter
- from typing import Any, Callable, Dict, List, Set, Tuple
- from isort.exceptions import (
- AssignmentsFormatMismatch,
- LiteralParsingFailure,
- LiteralSortTypeMismatch,
- )
- from isort.settings import DEFAULT_CONFIG, Config
- class ISortPrettyPrinter(PrettyPrinter):
- """an isort customized pretty printer for sorted literals"""
- def __init__(self, config: Config):
- super().__init__(width=config.line_length, compact=True)
- type_mapping: Dict[str, Tuple[type, Callable[[Any, ISortPrettyPrinter], str]]] = {}
- def assignments(code: str) -> str:
- values = {}
- for line in code.splitlines(keepends=True):
- if not line.strip():
- continue
- if " = " not in line:
- raise AssignmentsFormatMismatch(code)
- variable_name, value = line.split(" = ", 1)
- values[variable_name] = value
- return "".join(
- f"{variable_name} = {values[variable_name]}" for variable_name in sorted(values.keys())
- )
- def assignment(code: str, sort_type: str, extension: str, config: Config = DEFAULT_CONFIG) -> str:
- """Sorts the literal present within the provided code against the provided sort type,
- returning the sorted representation of the source code.
- """
- if sort_type == "assignments":
- return assignments(code)
- if sort_type not in type_mapping:
- raise ValueError(
- "Trying to sort using an undefined sort_type. "
- f"Defined sort types are {', '.join(type_mapping.keys())}."
- )
- variable_name, literal = code.split(" = ")
- variable_name = variable_name.lstrip()
- try:
- value = ast.literal_eval(literal)
- except Exception as error:
- raise LiteralParsingFailure(code, error)
- expected_type, sort_function = type_mapping[sort_type]
- if type(value) != expected_type:
- raise LiteralSortTypeMismatch(type(value), expected_type)
- printer = ISortPrettyPrinter(config)
- sorted_value_code = f"{variable_name} = {sort_function(value, printer)}"
- if config.formatting_function:
- sorted_value_code = config.formatting_function(
- sorted_value_code, extension, config
- ).rstrip()
- sorted_value_code += code[len(code.rstrip()) :]
- return sorted_value_code
- def register_type(
- name: str, kind: type
- ) -> Callable[[Callable[[Any, ISortPrettyPrinter], str]], Callable[[Any, ISortPrettyPrinter], str]]:
- """Registers a new literal sort type."""
- def wrap(
- function: Callable[[Any, ISortPrettyPrinter], str]
- ) -> Callable[[Any, ISortPrettyPrinter], str]:
- type_mapping[name] = (kind, function)
- return function
- return wrap
- @register_type("dict", dict)
- def _dict(value: Dict[Any, Any], printer: ISortPrettyPrinter) -> str:
- return printer.pformat(dict(sorted(value.items(), key=lambda item: item[1]))) # type: ignore
- @register_type("list", list)
- def _list(value: List[Any], printer: ISortPrettyPrinter) -> str:
- return printer.pformat(sorted(value))
- @register_type("unique-list", list)
- def _unique_list(value: List[Any], printer: ISortPrettyPrinter) -> str:
- return printer.pformat(list(sorted(set(value))))
- @register_type("set", set)
- def _set(value: Set[Any], printer: ISortPrettyPrinter) -> str:
- return "{" + printer.pformat(tuple(sorted(value)))[1:-1] + "}"
- @register_type("tuple", tuple)
- def _tuple(value: Tuple[Any, ...], printer: ISortPrettyPrinter) -> str:
- return printer.pformat(tuple(sorted(value)))
- @register_type("unique-tuple", tuple)
- def _unique_tuple(value: Tuple[Any, ...], printer: ISortPrettyPrinter) -> str:
- return printer.pformat(tuple(sorted(set(value))))
|