Improved the performance of the regular languages compiler: generate fewer and better regexes.

Author: jonathanslenders

prompt_toolkit/contrib/regular_languages/compiler.py

@@ -38,6 +38,7 @@
  m.variables().get('operator2') # Returns "add"
 
 """
+import itertools
 import re
 from typing import Callable, Dict, Iterable, Iterator, List
 from typing import Match as RegexMatch
@@ -175,10 +176,20 @@ def transform(node: Node) -> str:
 
  # `Repeat`.
  elif isinstance(node, Repeat):
- return "(?:%s){%i,%s}%s" % (
+ if node.max_repeat is None:
+ if node.min_repeat == 0:
+ repeat_sign = "*"
+ elif node.min_repeat == 1:
+ repeat_sign = "+"
+ else:
+ repeat_sign = "{%i,%s}" % (
+ node.min_repeat,
+ ("" if node.max_repeat is None else str(node.max_repeat)),
+ )
+
+ return "(?:%s)%s%s" % (
  transform(node.childnode),
- node.min_repeat,
- ("" if node.max_repeat is None else str(node.max_repeat)),
+ repeat_sign,
  ("" if node.greedy else "?"),
  )
  else:
@@ -207,26 +218,96 @@ def _transform_prefix(
  :param create_group_func: A callable which takes a `Node` and returns the next
  free name for this node.
  """
+ def contains_variable(node: Node) -> bool:
+ if isinstance(node, Regex):
+ return False
+ elif isinstance(node, Variable):
+ return True
+ elif isinstance(node, (Lookahead, Repeat)):
+ return contains_variable(node.childnode)
+ elif isinstance(node, (NodeSequence, AnyNode)):
+ return any(contains_variable(child) for child in node.children)
+
+ return False
 
  def transform(node: Node) -> Iterable[str]:
  # Generate regexes for all permutations of this OR. Each node
  # should be in front once.
  if isinstance(node, AnyNode):
+ # If we have a definition like:
+ # (?P<name> .*) | (?P<city> .*)
+ # Then we want to be able to generate completions for both the
+ # name as well as the city. We do this by yielding two
+ # different regular expressions, because the engine won't
+ # follow multiple paths, if multiple are possible.
+ children_with_variable = []
+ children_without_variable = []
  for c in node.children:
- for r in transform(c):
- yield "(?:%s)?" % r
+ if contains_variable(c):
+ children_with_variable.append(c)
+ else:
+ children_without_variable.append(c)
+
+ for c in children_with_variable:
+ yield from transform(c)
 
- # For a sequence. We can either have a match for the sequence
- # of all the children, or for an exact match of the first X
- # children, followed by a partial match of the next children.
+ # Merge options without variable together.
+ if children_without_variable:
+ yield "|".join(r for c in children_without_variable for r in transform(c))
+
+ # For a sequence. We can either have a match for the sequence of
+ # all the children, or for an exact match of the first X children,
+ # followed by a partial match of the next children.
+ # It is important here to yield a separate regex for each child
+ # that contains a variable. (With the variable at the end.)
  elif isinstance(node, NodeSequence):
+ # For all components in the sequence, compute prefix patterns,
+ # as well as full patterns.
+ complete = [
+ cls._transform(c, create_group_func) for c in node.children
+ ]
+ prefixes = [list(transform(c)) for c in node.children]
+ variable_nodes = [contains_variable(c) for c in node.children]
+
+ # However, if any child is contains a variable, we should
+ # yield an expression up to that point, so that we are sure
+ # this will be matched.
+ # (Otherwise, 'set (\s+ (?P<var> ..)) \s' won't complete "var".)
  for i in range(len(node.children)):
- a = [
- cls._transform(c, create_group_func) for c in node.children[:i]
- ]
-
- for c_str in transform(node.children[i]):
- yield "(?:%s)" % ("".join(a) + c_str)
+ if variable_nodes[i]:
+ for c_str in prefixes[i]:
+ yield "".join(complete[:i]) + c_str
+
+ # If there are non-variable nodes:
+ if not all(variable_nodes):
+ # If the input is: "[part1] [part2] [part3]", then this
+ # gets compiled into:
+ # (complete1 + (complete2 + (complete3 | partial3) | partial2) | partial1 )
+ # The partial matches in here can possibly contain many
+ # pattern (if it contained an `AnyNode`. We take the
+ # product of all combinations.)
+ # For nodes that contain a variable, we skip the "|partial" part here.
+ result = []
+
+ # Start with complete patterns.
+ for i, child in enumerate(node.children):
+ result.append(["(?:"])
+ result.append([complete[i]])
+
+ # Add prefix patterns.
+ for i, child in reversed(list(enumerate(node.children))):
+ if variable_nodes[i]:
+ # No need to yield a prefix for this one, we did
+ # the variable prefixes earlier.
+ result.append(')')
+ else:
+ result.append(["|(?:"])
+ # If this yields multiple, we should yield all combinations.
+ result.append(prefixes[i])
+ result.append(["))"])
+
+ for comb in itertools.product(*result):
+ yield "".join(comb)
 
  elif isinstance(node, Regex):
  yield "(?:%s)?" % node.regex
@@ -251,17 +332,20 @@ def transform(node: Node) -> Iterable[str]:
  # match, followed by a partial match.
  prefix = cls._transform(node.childnode, create_group_func)
 
- for c_str in transform(node.childnode):
- if node.max_repeat:
- repeat_sign = "{,%i}" % (node.max_repeat - 1)
- else:
- repeat_sign = "*"
- yield "(?:%s)%s%s(?:%s)?" % (
- prefix,
- repeat_sign,
- ("" if node.greedy else "?"),
- c_str,
- )
+ if node.max_repeat == 1:
+ yield from transform(node.childnode)
+ else:
+ for c_str in transform(node.childnode):
+ if node.max_repeat:
+ repeat_sign = "{,%i}" % (node.max_repeat - 1)
+ else:
+ repeat_sign = "*"
+ yield "(?:%s)%s%s%s" % (
+ prefix,
+ repeat_sign,
+ ("" if node.greedy else "?"),
+ c_str,
+ )
 
  else:
  raise TypeError("Got %r" % node)
@@ -343,7 +427,7 @@ def get_tuples() -> Iterable[Tuple[str, Tuple[int, int]]]:
 
  def _nodes_to_values(self) -> List[Tuple[str, str, Tuple[int, int]]]:
  """
- Returns list of list of (Node, string_value) tuples.
+ Returns list of (Node, string_value) tuples.
  """
 
  def is_none(sl: Tuple[int, int]) -> bool: