simplify a little the Answers parsing/formatting spaghetti code

clyngor/answers.py

@@ -8,6 +8,29 @@
 from clyngor import as_pyasp, parsing, utils
 
 
+def naive_parsing_of_answer_set(answer_set:str, *, discard_quotes:bool=False, parse_int:bool=True, parse_args:bool=True) -> [(str, tuple)]:
+    """Yield (pred, args), naively parsed from given answer set encoded as clingo output string"""
+    # print('NAIVE_PARSING_OF_ANSWER_SET:', answer_set, f'\t discard_quotes={discard_quotes}, parse_int={parse_int}, parse_args={parse_args}')
+    REG_ANSWER_SET = re.compile(r'([a-z][a-zA-Z0-9_]*)(\([^)]+\))?')
+
+    for match in REG_ANSWER_SET.finditer(answer_set):
+        pred, args = match.groups()
+        assert args is None or (args.startswith('(') and args.endswith(')'))
+        if args and parse_args:
+            args = args[1:-1]  # remove surrounding parens
+            if discard_quotes:
+                args = utils.remove_arguments_quotes(args)
+            args = tuple(
+                (int(arg) if parse_int and
+                 (arg[1:] if arg.startswith('-') else arg).isnumeric() else arg)
+                for arg in args.split(',')
+            ) if parse_args else args
+        elif args:  # args should not be parsed
+            args = args
+        yield pred, args or ()
+        # print('\t>', pred, args)
+
+
 class Answers:
     """Proxy to the solver, generated by solving methods like solve.solve
     or inline.ASP.
@@ -165,7 +188,6 @@ def no_arg(self):
     def __next__(self):
         return next(iter(self))
 
-
     def __iter__(self):
         """Yield answer sets"""
         for answer_set, optimization, optimality in self._answers:
@@ -185,46 +207,37 @@ def clean_resources(self):
 
 
     def _parse_answer(self, answer_set:str) -> iter:
-        """Yield atoms as (pred, args) according to parsing options"""
-        REG_ANSWER_SET = re.compile(r'([a-z][a-zA-Z0-9_]*)(\([^)]+\))?')
-        if self._careful_parsing:
+        """Yield atoms as (pred, args) from given answer set"""
+        if isinstance(answer_set, str) and self._careful_parsing:
+            # print('CAREFUL PARSING:', answer_set)
             # _discard_quotes is incompatible with atoms_as_string and as_pyasp.
             # atom_as_string: remove the quotes delimiting arguments.
             # as_pyasp: remove the quotes for the arguments.
             yield from parsing.Parser(
                 self._collapse_atoms, self._collapse_args,
-                self._discard_quotes and not self._atoms_as_string,
+                self._discard_quotes and not self._collapse_atoms,
                 self._first_arg_only,
                 parse_integer=self._parse_int
             ).parse_terms(answer_set)
-
-        else:  # the good ol' split
-            current_answer = set()
-            for match in REG_ANSWER_SET.finditer(answer_set):
-                if self._atoms_as_string:
-                    pred, args = match.groups()
-                    if args is None:  # atom with no arg
-                        args = ''
-                    elif self._first_arg_only:
-                        args = args.split(',')[0]
-                        if not args.endswith(')'):
-                            args = args + ')'
-                    yield pred + args  # just send the match
-                    continue
-                pred, args = match.groups()
-                assert args is None or (args.startswith('(') and args.endswith(')'))
-                if args:
-                    args = args[1:-1]
-                    if self._discard_quotes:
-                        args = utils.remove_arguments_quotes(args)
-                    if not self._collapse_atoms:  # else: atom as string
-                        # parse also integers, if asked to
-                        args = tuple(
-                            (int(arg) if self._parse_int and
-                             (arg[1:] if arg.startswith('-') else arg).isnumeric() else arg)
-                            for arg in args.split(',')
-                        )
-                yield pred, args or ()
+        elif isinstance(answer_set, str):  # the good ol' split
+            # print('THE GOOD OLD SPLIT:', f"discard_quotes={self._discard_quotes}  collapse_atoms={self._collapse_atoms}")
+            yield from self.__finish_parsing(naive_parsing_of_answer_set(answer_set, discard_quotes=self._discard_quotes and not self._collapse_atoms, parse_int=self._parse_int, parse_args=True or self._collapse_args or self._first_arg_only))
+        elif isinstance(answer_set, (set, tuple)) and all(isinstance(atom, tuple) for atom in answer_set):  # already parsed
+            # print('FROM SET OR TUPLE')
+            yield from self.__finish_parsing(answer_set)
+        else:  # unknown format
+            raise ValueError(f"unknow answer set format: {type(answer_set)}, {answer_set}")
+
+    def __finish_parsing(self, answer_set:[(str, tuple)]) -> [(str, tuple) or str]:
+        """Modify (pred, args) atoms according to parsing options"""
+        for pred, args in answer_set:
+            if self._first_arg_only:
+                args = [args[0]] if args else args
+            if self._atoms_as_string:
+                args = ('(' + ','.join(map(str, args)) + ')') if args else ''
+                yield pred + args
+            else:
+                yield pred, args
 
 
     def _format(self, answer_set) -> dict or frozenset: