Add attritemgetter partials

Resolves #787.

.pre-commit-config.yaml

@@ -24,13 +24,13 @@ repos:
         args:
         - --autofix
 -   repo: https://github.com/pycqa/flake8
-    rev: 6.0.0
+    rev: 6.1.0
     hooks:
     -   id: flake8
         args:
         - --ignore=W503,E501,E265,E402,F405,E305,E126
 -   repo: https://github.com/pre-commit/mirrors-autopep8
-    rev: v2.0.2
+    rev: v2.0.4
     hooks:
     -   id: autopep8
         args:

__coconut__/__init__.pyi

@@ -640,6 +640,12 @@ def _coconut_iter_getitem(
     ...
 
 
+def _coconut_attritemgetter(
+    attr: _t.Optional[_t.Text],
+    *is_iter_and_items: _t.Tuple[_t.Tuple[bool, _t.Any], ...],
+) -> _t.Callable[[_t.Any], _t.Any]: ...
+
+
 def _coconut_base_compose(
     func: _t.Callable[[_T], _t.Any],
     *func_infos: _t.Tuple[_Callable, int, bool],

coconut/__coconut__.pyi

@@ -1,2 +1,2 @@
 from __coconut__ import *
-from __coconut__ import _coconut_tail_call, _coconut_tco, _coconut_call_set_names, _coconut_handle_cls_kwargs, _coconut_handle_cls_stargs, _namedtuple_of, _coconut, _coconut_Expected, _coconut_MatchError, _coconut_SupportsAdd, _coconut_SupportsMinus, _coconut_SupportsMul, _coconut_SupportsPow, _coconut_SupportsTruediv, _coconut_SupportsFloordiv, _coconut_SupportsMod, _coconut_SupportsAnd, _coconut_SupportsXor, _coconut_SupportsOr, _coconut_SupportsLshift, _coconut_SupportsRshift, _coconut_SupportsMatmul, _coconut_SupportsInv, _coconut_Expected, _coconut_MatchError, _coconut_iter_getitem, _coconut_base_compose, _coconut_forward_compose, _coconut_back_compose, _coconut_forward_star_compose, _coconut_back_star_compose, _coconut_forward_dubstar_compose, _coconut_back_dubstar_compose, _coconut_pipe, _coconut_star_pipe, _coconut_dubstar_pipe, _coconut_back_pipe, _coconut_back_star_pipe, _coconut_back_dubstar_pipe, _coconut_none_pipe, _coconut_none_star_pipe, _coconut_none_dubstar_pipe, _coconut_bool_and, _coconut_bool_or, _coconut_none_coalesce, _coconut_minus, _coconut_map, _coconut_partial, _coconut_get_function_match_error, _coconut_base_pattern_func, _coconut_addpattern, _coconut_sentinel, _coconut_assert, _coconut_raise, _coconut_mark_as_match, _coconut_reiterable, _coconut_self_match_types, _coconut_dict_merge, _coconut_exec, _coconut_comma_op, _coconut_multi_dim_arr, _coconut_mk_anon_namedtuple, _coconut_matmul, _coconut_py_str, _coconut_flatten, _coconut_multiset, _coconut_back_none_pipe, _coconut_back_none_star_pipe, _coconut_back_none_dubstar_pipe, _coconut_forward_none_compose, _coconut_back_none_compose, _coconut_forward_none_star_compose, _coconut_back_none_star_compose, _coconut_forward_none_dubstar_compose, _coconut_back_none_dubstar_compose, _coconut_call_or_coefficient, _coconut_in, _coconut_not_in
+from __coconut__ import _coconut_tail_call, _coconut_tco, _coconut_call_set_names, _coconut_handle_cls_kwargs, _coconut_handle_cls_stargs, _namedtuple_of, _coconut, _coconut_Expected, _coconut_MatchError, _coconut_SupportsAdd, _coconut_SupportsMinus, _coconut_SupportsMul, _coconut_SupportsPow, _coconut_SupportsTruediv, _coconut_SupportsFloordiv, _coconut_SupportsMod, _coconut_SupportsAnd, _coconut_SupportsXor, _coconut_SupportsOr, _coconut_SupportsLshift, _coconut_SupportsRshift, _coconut_SupportsMatmul, _coconut_SupportsInv, _coconut_Expected, _coconut_MatchError, _coconut_iter_getitem, _coconut_base_compose, _coconut_forward_compose, _coconut_back_compose, _coconut_forward_star_compose, _coconut_back_star_compose, _coconut_forward_dubstar_compose, _coconut_back_dubstar_compose, _coconut_pipe, _coconut_star_pipe, _coconut_dubstar_pipe, _coconut_back_pipe, _coconut_back_star_pipe, _coconut_back_dubstar_pipe, _coconut_none_pipe, _coconut_none_star_pipe, _coconut_none_dubstar_pipe, _coconut_bool_and, _coconut_bool_or, _coconut_none_coalesce, _coconut_minus, _coconut_map, _coconut_partial, _coconut_get_function_match_error, _coconut_base_pattern_func, _coconut_addpattern, _coconut_sentinel, _coconut_assert, _coconut_raise, _coconut_mark_as_match, _coconut_reiterable, _coconut_self_match_types, _coconut_dict_merge, _coconut_exec, _coconut_comma_op, _coconut_multi_dim_arr, _coconut_mk_anon_namedtuple, _coconut_matmul, _coconut_py_str, _coconut_flatten, _coconut_multiset, _coconut_back_none_pipe, _coconut_back_none_star_pipe, _coconut_back_none_dubstar_pipe, _coconut_forward_none_compose, _coconut_back_none_compose, _coconut_forward_none_star_compose, _coconut_back_none_star_compose, _coconut_forward_none_dubstar_compose, _coconut_back_none_dubstar_compose, _coconut_call_or_coefficient, _coconut_in, _coconut_not_in, _coconut_attritemgetter

coconut/compiler/compiler.py

@@ -2747,7 +2747,7 @@ def pipe_item_split(self, tokens, loc):
             - (expr,) for expression
             - (func, pos_args, kwd_args) for partial
             - (name, args) for attr/method
-            - (op, args)+ for itemgetter
+            - (attr, [(op, args)]) for itemgetter
             - (op, arg) for right op partial
         """
         # list implies artificial tokens, which must be expr
@@ -2762,8 +2762,12 @@ def pipe_item_split(self, tokens, loc):
             name, args = attrgetter_atom_split(tokens)
             return "attrgetter", (name, args)
         elif "itemgetter" in tokens:
-            internal_assert(len(tokens) >= 2, "invalid itemgetter pipe item tokens", tokens)
-            return "itemgetter", tokens
+            if len(tokens) == 1:
+                attr = None
+                ops_and_args, = tokens
+            else:
+                attr, ops_and_args = tokens
+            return "itemgetter", (attr, ops_and_args)
         elif "op partial" in tokens:
             inner_toks, = tokens
             if "left partial" in inner_toks:
@@ -2853,12 +2857,13 @@ def pipe_handle(self, original, loc, tokens, **kwargs):
                 elif name == "itemgetter":
                     if stars:
                         raise CoconutDeferredSyntaxError("cannot star pipe into item getting", loc)
-                    self.internal_assert(len(split_item) % 2 == 0, original, loc, "invalid itemgetter pipe tokens", split_item)
-                    out = subexpr
-                    for i in range(0, len(split_item), 2):
-                        op, args = split_item[i:i + 2]
+                    attr, ops_and_args = split_item
+                    out = "(" + subexpr + ")"
+                    if attr is not None:
+                        out += "." + attr
+                    for op, args in ops_and_args:
                         if op == "[":
-                            fmtstr = "({x})[{args}]"
+                            fmtstr = "{x}[{args}]"
                         elif op == "$[":
                             fmtstr = "_coconut_iter_getitem({x}, ({args}))"
                         else:

coconut/compiler/grammar.py

@@ -437,22 +437,24 @@ def subscriptgroup_handle(tokens):
 
 def itemgetter_handle(tokens):
     """Process implicit itemgetter partials."""
-    if len(tokens) == 2:
-        op, args = tokens
+    if len(tokens) == 1:
+        attr = None
+        ops_and_args, = tokens
+    else:
+        attr, ops_and_args = tokens
+    if attr is None and len(ops_and_args) == 1:
+        (op, args), = ops_and_args
         if op == "[":
             return "_coconut.operator.itemgetter((" + args + "))"
         elif op == "$[":
             return "_coconut.functools.partial(_coconut_iter_getitem, index=(" + args + "))"
         else:
             raise CoconutInternalException("invalid implicit itemgetter type", op)
-    elif len(tokens) > 2:
-        internal_assert(len(tokens) % 2 == 0, "invalid itemgetter composition tokens", tokens)
-        itemgetters = []
-        for i in range(0, len(tokens), 2):
-            itemgetters.append(itemgetter_handle(tokens[i:i + 2]))
-        return "_coconut_forward_compose(" + ", ".join(itemgetters) + ")"
     else:
-        raise CoconutInternalException("invalid implicit itemgetter tokens", tokens)
+        return "_coconut_attritemgetter({attr}, {is_iter_and_items})".format(
+            attr=repr(attr),
+            is_iter_and_items=", ".join("({is_iter}, ({item}))".format(is_iter=op == "$[", item=args) for op, args in ops_and_args),
+        )
 
 
 def class_suite_handle(tokens):
@@ -1300,11 +1302,21 @@ class Grammar(object):
         lparen + Optional(methodcaller_args) + rparen.suppress()
     )
     attrgetter_atom = attach(attrgetter_atom_tokens, attrgetter_atom_handle)
-    itemgetter_atom_tokens = dot.suppress() + OneOrMore(condense(Optional(dollar) + lbrack) + subscriptgrouplist + rbrack.suppress())
+
+    itemgetter_atom_tokens = (
+        dot.suppress()
+        + Optional(unsafe_dotted_name)
+        + Group(OneOrMore(Group(
+            condense(Optional(dollar) + lbrack)
+            + subscriptgrouplist
+            + rbrack.suppress()
+        )))
+    )
     itemgetter_atom = attach(itemgetter_atom_tokens, itemgetter_handle)
+
     implicit_partial_atom = (
-        attrgetter_atom
-        | itemgetter_atom
+        itemgetter_atom
+        | attrgetter_atom
         | fixto(dot + lbrack + rbrack, "_coconut.operator.getitem")
         | fixto(dot + dollar + lbrack + rbrack, "_coconut_iter_getitem")
     )
@@ -1485,8 +1497,8 @@ class Grammar(object):
     pipe_item = (
         # we need the pipe_op since any of the atoms could otherwise be the start of an expression
         labeled_group(keyword("await"), "await") + pipe_op
-        | labeled_group(attrgetter_atom_tokens, "attrgetter") + pipe_op
         | labeled_group(itemgetter_atom_tokens, "itemgetter") + pipe_op
+        | labeled_group(attrgetter_atom_tokens, "attrgetter") + pipe_op
         | labeled_group(partial_atom_tokens, "partial") + pipe_op
         | labeled_group(partial_op_atom_tokens, "op partial") + pipe_op
         # expr must come at end
@@ -1495,8 +1507,8 @@ class Grammar(object):
     pipe_augassign_item = (
         # should match pipe_item but with pipe_op -> end_simple_stmt_item and no expr
         labeled_group(keyword("await"), "await") + end_simple_stmt_item
-        | labeled_group(attrgetter_atom_tokens, "attrgetter") + end_simple_stmt_item
         | labeled_group(itemgetter_atom_tokens, "itemgetter") + end_simple_stmt_item
+        | labeled_group(attrgetter_atom_tokens, "attrgetter") + end_simple_stmt_item
         | labeled_group(partial_atom_tokens, "partial") + end_simple_stmt_item
         | labeled_group(partial_op_atom_tokens, "op partial") + end_simple_stmt_item
     )
@@ -1505,8 +1517,8 @@ class Grammar(object):
         # we need longest here because there's no following pipe_op we can use as above
         | longest(
             keyword("await")("await"),
-            attrgetter_atom_tokens("attrgetter"),
             itemgetter_atom_tokens("itemgetter"),
+            attrgetter_atom_tokens("attrgetter"),
             partial_atom_tokens("partial"),
             partial_op_atom_tokens("op partial"),
             comp_pipe_expr("expr"),

coconut/compiler/header.py

@@ -586,7 +586,7 @@ async def __anext__(self):
     #  (extra_format_dict is to keep indentation levels matching)
     extra_format_dict = dict(
         # when anything is added to this list it must also be added to *both* __coconut__ stub files
-        underscore_imports="{tco_comma}{call_set_names_comma}{handle_cls_args_comma}_namedtuple_of, _coconut, _coconut_Expected, _coconut_MatchError, _coconut_SupportsAdd, _coconut_SupportsMinus, _coconut_SupportsMul, _coconut_SupportsPow, _coconut_SupportsTruediv, _coconut_SupportsFloordiv, _coconut_SupportsMod, _coconut_SupportsAnd, _coconut_SupportsXor, _coconut_SupportsOr, _coconut_SupportsLshift, _coconut_SupportsRshift, _coconut_SupportsMatmul, _coconut_SupportsInv, _coconut_iter_getitem, _coconut_base_compose, _coconut_forward_compose, _coconut_back_compose, _coconut_forward_star_compose, _coconut_back_star_compose, _coconut_forward_dubstar_compose, _coconut_back_dubstar_compose, _coconut_pipe, _coconut_star_pipe, _coconut_dubstar_pipe, _coconut_back_pipe, _coconut_back_star_pipe, _coconut_back_dubstar_pipe, _coconut_none_pipe, _coconut_none_star_pipe, _coconut_none_dubstar_pipe, _coconut_bool_and, _coconut_bool_or, _coconut_none_coalesce, _coconut_minus, _coconut_map, _coconut_partial, _coconut_get_function_match_error, _coconut_base_pattern_func, _coconut_addpattern, _coconut_sentinel, _coconut_assert, _coconut_raise, _coconut_mark_as_match, _coconut_reiterable, _coconut_self_match_types, _coconut_dict_merge, _coconut_exec, _coconut_comma_op, _coconut_multi_dim_arr, _coconut_mk_anon_namedtuple, _coconut_matmul, _coconut_py_str, _coconut_flatten, _coconut_multiset, _coconut_back_none_pipe, _coconut_back_none_star_pipe, _coconut_back_none_dubstar_pipe, _coconut_forward_none_compose, _coconut_back_none_compose, _coconut_forward_none_star_compose, _coconut_back_none_star_compose, _coconut_forward_none_dubstar_compose, _coconut_back_none_dubstar_compose, _coconut_call_or_coefficient, _coconut_in, _coconut_not_in".format(**format_dict),
+        underscore_imports="{tco_comma}{call_set_names_comma}{handle_cls_args_comma}_namedtuple_of, _coconut, _coconut_Expected, _coconut_MatchError, _coconut_SupportsAdd, _coconut_SupportsMinus, _coconut_SupportsMul, _coconut_SupportsPow, _coconut_SupportsTruediv, _coconut_SupportsFloordiv, _coconut_SupportsMod, _coconut_SupportsAnd, _coconut_SupportsXor, _coconut_SupportsOr, _coconut_SupportsLshift, _coconut_SupportsRshift, _coconut_SupportsMatmul, _coconut_SupportsInv, _coconut_iter_getitem, _coconut_base_compose, _coconut_forward_compose, _coconut_back_compose, _coconut_forward_star_compose, _coconut_back_star_compose, _coconut_forward_dubstar_compose, _coconut_back_dubstar_compose, _coconut_pipe, _coconut_star_pipe, _coconut_dubstar_pipe, _coconut_back_pipe, _coconut_back_star_pipe, _coconut_back_dubstar_pipe, _coconut_none_pipe, _coconut_none_star_pipe, _coconut_none_dubstar_pipe, _coconut_bool_and, _coconut_bool_or, _coconut_none_coalesce, _coconut_minus, _coconut_map, _coconut_partial, _coconut_get_function_match_error, _coconut_base_pattern_func, _coconut_addpattern, _coconut_sentinel, _coconut_assert, _coconut_raise, _coconut_mark_as_match, _coconut_reiterable, _coconut_self_match_types, _coconut_dict_merge, _coconut_exec, _coconut_comma_op, _coconut_multi_dim_arr, _coconut_mk_anon_namedtuple, _coconut_matmul, _coconut_py_str, _coconut_flatten, _coconut_multiset, _coconut_back_none_pipe, _coconut_back_none_star_pipe, _coconut_back_none_dubstar_pipe, _coconut_forward_none_compose, _coconut_back_none_compose, _coconut_forward_none_star_compose, _coconut_back_none_star_compose, _coconut_forward_none_dubstar_compose, _coconut_back_none_dubstar_compose, _coconut_call_or_coefficient, _coconut_in, _coconut_not_in, _coconut_attritemgetter".format(**format_dict),
         import_typing=pycondition(
             (3, 5),
             if_ge='''