|
57 | 57 | Hashable, |
58 | 58 | Iterable, |
59 | 59 | Iterator, |
| 60 | + Literal, |
60 | 61 | Mapping, |
61 | 62 | MutableMapping, |
62 | 63 | MutableSet, |
| 64 | + Sequence, |
63 | 65 | TypeVar, |
64 | 66 | cast, |
65 | 67 | overload, |
|
69 | 71 | import pandas as pd |
70 | 72 |
|
71 | 73 | if TYPE_CHECKING: |
72 | | - from .types import ErrorOptionsWithWarn |
| 74 | + from .types import Dims, ErrorOptionsWithWarn, OrderedDims |
73 | 75 |
|
74 | 76 | K = TypeVar("K") |
75 | 77 | V = TypeVar("V") |
@@ -887,15 +889,17 @@ def drop_dims_from_indexers( |
887 | 889 |
|
888 | 890 |
|
889 | 891 | def drop_missing_dims( |
890 | | - supplied_dims: Collection, dims: Collection, missing_dims: ErrorOptionsWithWarn |
891 | | -) -> Collection: |
| 892 | + supplied_dims: Iterable[Hashable], |
| 893 | + dims: Iterable[Hashable], |
| 894 | + missing_dims: ErrorOptionsWithWarn, |
| 895 | +) -> Iterable[Hashable]: |
892 | 896 | """Depending on the setting of missing_dims, drop any dimensions from supplied_dims that |
893 | 897 | are not present in dims. |
894 | 898 |
|
895 | 899 | Parameters |
896 | 900 | ---------- |
897 | | - supplied_dims : dict |
898 | | - dims : sequence |
| 901 | + supplied_dims : Iterable of Hashable |
| 902 | + dims : Iterable of Hashable |
899 | 903 | missing_dims : {"raise", "warn", "ignore"} |
900 | 904 | """ |
901 | 905 |
|
@@ -928,6 +932,158 @@ def drop_missing_dims( |
928 | 932 | ) |
929 | 933 |
|
930 | 934 |
|
| 935 | +T_None = TypeVar("T_None", None, "ellipsis") |
| 936 | + |
| 937 | + |
| 938 | +@overload |
| 939 | +def parse_dims( |
| 940 | + dim: str | Iterable[Hashable] | T_None, |
| 941 | + all_dims: tuple[Hashable, ...], |
| 942 | + *, |
| 943 | + check_exists: bool = True, |
| 944 | + replace_none: Literal[True] = True, |
| 945 | +) -> tuple[Hashable, ...]: |
| 946 | + ... |
| 947 | + |
| 948 | + |
| 949 | +@overload |
| 950 | +def parse_dims( |
| 951 | + dim: str | Iterable[Hashable] | T_None, |
| 952 | + all_dims: tuple[Hashable, ...], |
| 953 | + *, |
| 954 | + check_exists: bool = True, |
| 955 | + replace_none: Literal[False], |
| 956 | +) -> tuple[Hashable, ...] | T_None: |
| 957 | + ... |
| 958 | + |
| 959 | + |
| 960 | +def parse_dims( |
| 961 | + dim: Dims, |
| 962 | + all_dims: tuple[Hashable, ...], |
| 963 | + *, |
| 964 | + check_exists: bool = True, |
| 965 | + replace_none: bool = True, |
| 966 | +) -> tuple[Hashable, ...] | None | ellipsis: |
| 967 | + """Parse one or more dimensions. |
| 968 | +
|
| 969 | + A single dimension must be always a str, multiple dimensions |
| 970 | + can be Hashables. This supports e.g. using a tuple as a dimension. |
| 971 | + If you supply e.g. a set of dimensions the order cannot be |
| 972 | + conserved, but for sequences it will be. |
| 973 | +
|
| 974 | + Parameters |
| 975 | + ---------- |
| 976 | + dim : str, Iterable of Hashable, "..." or None |
| 977 | + Dimension(s) to parse. |
| 978 | + all_dims : tuple of Hashable |
| 979 | + All possible dimensions. |
| 980 | + check_exists: bool, default: True |
| 981 | + if True, check if dim is a subset of all_dims. |
| 982 | + replace_none : bool, default: True |
| 983 | + If True, return all_dims if dim is None or "...". |
| 984 | +
|
| 985 | + Returns |
| 986 | + ------- |
| 987 | + parsed_dims : tuple of Hashable |
| 988 | + Input dimensions as a tuple. |
| 989 | + """ |
| 990 | + if dim is None or dim is ...: |
| 991 | + if replace_none: |
| 992 | + return all_dims |
| 993 | + return dim |
| 994 | + if isinstance(dim, str): |
| 995 | + dim = (dim,) |
| 996 | + if check_exists: |
| 997 | + _check_dims(set(dim), set(all_dims)) |
| 998 | + return tuple(dim) |
| 999 | + |
| 1000 | + |
| 1001 | +@overload |
| 1002 | +def parse_ordered_dims( |
| 1003 | + dim: str | Sequence[Hashable | ellipsis] | T_None, |
| 1004 | + all_dims: tuple[Hashable, ...], |
| 1005 | + *, |
| 1006 | + check_exists: bool = True, |
| 1007 | + replace_none: Literal[True] = True, |
| 1008 | +) -> tuple[Hashable, ...]: |
| 1009 | + ... |
| 1010 | + |
| 1011 | + |
| 1012 | +@overload |
| 1013 | +def parse_ordered_dims( |
| 1014 | + dim: str | Sequence[Hashable | ellipsis] | T_None, |
| 1015 | + all_dims: tuple[Hashable, ...], |
| 1016 | + *, |
| 1017 | + check_exists: bool = True, |
| 1018 | + replace_none: Literal[False], |
| 1019 | +) -> tuple[Hashable, ...] | T_None: |
| 1020 | + ... |
| 1021 | + |
| 1022 | + |
| 1023 | +def parse_ordered_dims( |
| 1024 | + dim: OrderedDims, |
| 1025 | + all_dims: tuple[Hashable, ...], |
| 1026 | + *, |
| 1027 | + check_exists: bool = True, |
| 1028 | + replace_none: bool = True, |
| 1029 | +) -> tuple[Hashable, ...] | None | ellipsis: |
| 1030 | + """Parse one or more dimensions. |
| 1031 | +
|
| 1032 | + A single dimension must be always a str, multiple dimensions |
| 1033 | + can be Hashables. This supports e.g. using a tuple as a dimension. |
| 1034 | + An ellipsis ("...") in a sequence of dimensions will be |
| 1035 | + replaced with all remaining dimensions. This only makes sense when |
| 1036 | + the input is a sequence and not e.g. a set. |
| 1037 | +
|
| 1038 | + Parameters |
| 1039 | + ---------- |
| 1040 | + dim : str, Sequence of Hashable or "...", "..." or None |
| 1041 | + Dimension(s) to parse. If "..." appears in a Sequence |
| 1042 | + it always gets replaced with all remaining dims |
| 1043 | + all_dims : tuple of Hashable |
| 1044 | + All possible dimensions. |
| 1045 | + check_exists: bool, default: True |
| 1046 | + if True, check if dim is a subset of all_dims. |
| 1047 | + replace_none : bool, default: True |
| 1048 | + If True, return all_dims if dim is None. |
| 1049 | +
|
| 1050 | + Returns |
| 1051 | + ------- |
| 1052 | + parsed_dims : tuple of Hashable |
| 1053 | + Input dimensions as a tuple. |
| 1054 | + """ |
| 1055 | + if dim is not None and dim is not ... and not isinstance(dim, str) and ... in dim: |
| 1056 | + dims_set: set[Hashable | ellipsis] = set(dim) |
| 1057 | + all_dims_set = set(all_dims) |
| 1058 | + if check_exists: |
| 1059 | + _check_dims(dims_set, all_dims_set) |
| 1060 | + if len(all_dims_set) != len(all_dims): |
| 1061 | + raise ValueError("Cannot use ellipsis with repeated dims") |
| 1062 | + dims = tuple(dim) |
| 1063 | + if dims.count(...) > 1: |
| 1064 | + raise ValueError("More than one ellipsis supplied") |
| 1065 | + other_dims = tuple(d for d in all_dims if d not in dims_set) |
| 1066 | + idx = dims.index(...) |
| 1067 | + return dims[:idx] + other_dims + dims[idx + 1 :] |
| 1068 | + else: |
| 1069 | + # mypy cannot resolve that the sequence cannot contain "..." |
| 1070 | + return parse_dims( # type: ignore[call-overload] |
| 1071 | + dim=dim, |
| 1072 | + all_dims=all_dims, |
| 1073 | + check_exists=check_exists, |
| 1074 | + replace_none=replace_none, |
| 1075 | + ) |
| 1076 | + |
| 1077 | + |
| 1078 | +def _check_dims(dim: set[Hashable | ellipsis], all_dims: set[Hashable]) -> None: |
| 1079 | + wrong_dims = dim - all_dims |
| 1080 | + if wrong_dims and wrong_dims != {...}: |
| 1081 | + wrong_dims_str = ", ".join(f"'{d!s}'" for d in wrong_dims) |
| 1082 | + raise ValueError( |
| 1083 | + f"Dimension(s) {wrong_dims_str} do not exist. Expected one or more of {all_dims}" |
| 1084 | + ) |
| 1085 | + |
| 1086 | + |
931 | 1087 | _Accessor = TypeVar("_Accessor") |
932 | 1088 |
|
933 | 1089 |
|
|
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=40){kind=avatar}
0 commit comments