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.strip() literal = literal.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))))