Assignments/redirects identified during parsing

We were half-parsing these (and redirects) during the
construction of the Command object. Do this during
parsing instead; it's required to identify heredocs.

Author: jan-g

psh/model.py

@@ -1,3 +1,4 @@
+import attr
 import contextlib
 import fcntl
 import io
@@ -106,6 +107,18 @@ def __getitem__(self, key):
             return super().__getitem__(key)
 
 
+@attr.s
+class Assignment:
+    var = attr.ib()
+    expr = attr.ib()
+
+    @staticmethod
+    def run(env, assignments):
+        for a in assignments:
+            assert isinstance(a.var, str)
+            env[str(a.var)] = a.expr.evaluate(env)
+
+
 class ConstantString(Comparable):
     """ An uninterpreted piece of string """
     def __init__(self, s, *args, **kwargs):
@@ -350,27 +363,9 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # Extract assignments and redirects
         self.assignments = []
-        items = []
-        might_be_assignment = True
-
-        for item in self:
-            if isinstance(item, Word):
-                if might_be_assignment and item.matches_assignment():
-                    self.with_assignment(item)
-                    continue
-                might_be_assignment = False
-
-                redirect = Redirect.from_word(item)
-                if redirect is not None:
-                    self.with_redirect(redirect)
-                    continue
-
-            items.append(item)
-
-        self[:] = items
 
-    def with_assignment(self, assignment):
-        self.assignments.append(assignment)
+    def with_assignment(self, *assignment):
+        self.assignments.extend(assignment)
         return self
 
     def is_null(self):
@@ -379,9 +374,7 @@ def is_null(self):
     def execute(self, env, input=None, output=None, error=None):
         LOG.debug("executing: %s", self)
         assert env.permit_execution
-        for var, _, *rest in self.assignments:
-            assert isinstance(var, Id)
-            env[str(var)] = Word(rest).evaluate(env)
+        Assignment.run(env, self.assignments)
         if env.permit_execution and len(self) > 0:
             args = [item.evaluate(env) for item in self]
             if args[0] in env.builtins:

psh/parser.py

@@ -1,32 +1,57 @@
 from functools import partial
-from parsy import eof, regex, generate, string, whitespace, ParseError, fail, seq, success, string_from, eof, any_char
-from .model import (ConstantString, Token, Id, VarRef, Word, Arith,
+from parsy import eof, regex, generate, string, ParseError, fail, seq, success, string_from, eof, any_char
+from parsy_extn import monkeypatch_parsy
+
+from .model import (ConstantString, Token, Id, VarRef, Word, Arith, Assignment,
                     Command, CommandSequence, CommandPipe, While, If, Function,
-                    RedirectFrom, RedirectTo, RedirectDup,
+                    Redirect, RedirectFrom, RedirectTo, RedirectDup,
                     MaybeDoubleQuoted,)
 
 
+# All our parser use may need to carry notes forward.
+# This cost is only paid if we actually muck around with heredocs.
+monkeypatch_parsy()
+
+
 ws = regex('([ \t]|\\\\\n)+')
 eol = string('\n')
 
+whitespace = ws | eol
+
+
 # End of statement
 EOF = object()
-eos = ws.optional() >> (regex('[;\n]') | eof.result(EOF))
+eos = ws.optional() >> (string(";") | eol | eof.result(EOF))
 
 @generate("command")
 def command():
     words = []
+    assignments = []
+    redirs = []
+    assignments_possible = True
     while True:
         yield ws.optional()
-        w = yield word
+        if assignments_possible:
+            w = yield assignment | redirect | word
+        else:
+            w = yield redirect | word
+        if isinstance(w, Word):
+            assignments_possible = False
+        elif isinstance(w, Redirect):
+            redirs.append(w)
+            continue
+        elif isinstance(w, Assignment):
+            assignments.append(w)
+            continue
+
         if not w:
             break
         if len(words) == 0 and w.matches_reserved("while", "do", "done", "if", "then", "elif", "else", "fi"):
             return fail("can't have a reserved word here")
 
         words.append(w)
 
-    cmd = Command(words)
+    cmd = Command(words).with_assignment(*assignments).with_redirect(*redirs)
     return cmd
 
 
@@ -224,6 +249,7 @@ def expr_add():
     lambda x: x[0] if len(x) == 1 and isinstance(x[0], Word) else
     Word([i for i in x if i != Token("")]))
 
+assignment = seq(variable_id, string("="), word).map(lambda vew: Assignment(vew[0], vew[2]))
 
 redirect_dup_from_n = seq(regex("[0-9]+"), string("<&") >> word).combine(RedirectDup)
 redirect_dup_from = (string("<&") >> word).map(partial(RedirectDup, 0))

setup.py

@@ -25,6 +25,7 @@ def read_file(fn):
 
     install_requires=[
                       "argcomplete",
+                      "attrs",
                       "parsy",
                       "parsy-extn",
                       "prompt_toolkit",

test/test_parser.py

@@ -2,7 +2,7 @@
 import pytest
 
 from psh.parser import command, command_sequence
-from psh.model import Word, ConstantString, Command, VarRef, Id, Token, CommandSequence, CommandPipe, While, If
+from psh.model import Word, ConstantString, Assignment, Command, VarRef, Id, Token, CommandSequence, CommandPipe, While, If
 from psh.local import make_env
 
 
@@ -18,11 +18,15 @@
         ("$(cat $foo $bar)", Command([Word([CommandSequence([Command([Word([ConstantString("cat")]),
                                                                       Word([VarRef("foo")]),
                                                                       Word([VarRef("bar")])])])])])),
-        ("a=2", Command([]).with_assignment(Word([Id("a"), Token("="), ConstantString("2")]))),
+        ("a=2", Command([]).with_assignment(Assignment("a", Word([ConstantString("2")])))),
         ("a=1 b=2 echo $a$b", Command([Word([Id("echo")]),
                                        Word([VarRef("a"), VarRef("b")])]).
-         with_assignment(Word([Id("a"), Token("="), ConstantString("1")])).
-         with_assignment(Word([Id("b"), Token("="), ConstantString("2")]))),
+         with_assignment(Assignment("a", Word([ConstantString("1")])),
+                         Assignment("b", Word([ConstantString("2")])))),
+        ("a=2 echo b=1", Command([Word([Id("echo")]),
+                                  Word([Id("b"), Token("="), ConstantString("1")])]).
+         with_assignment(Assignment("a", Word([ConstantString("2")])))),
+
 ))
 def test_basic(text, expected):
     cmd = command.parse(text)