From 748c8eebddc63edadfe50648d4d9f7b2a0b6ff9b Mon Sep 17 00:00:00 2001
From: Stephan Zwicknagl <stephanzwicknagl@gmx.de>
Date: Wed, 20 Dec 2023 18:02:21 -0700
Subject: [PATCH] potassco/viasp#40: improve statement visiting

---
 backend/src/viasp/__main__.py  |  2 ++
 backend/src/viasp/asp/reify.py | 34 ++++++++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/backend/src/viasp/__main__.py b/backend/src/viasp/__main__.py
index 55c10df3..401b5ed6 100644
--- a/backend/src/viasp/__main__.py
+++ b/backend/src/viasp/__main__.py
@@ -2,6 +2,7 @@
 import textwrap
 import webbrowser
 import importlib.metadata
+from clingo.script import enable_python
 
 from viasp import Control
 from viasp.server import startup
@@ -84,6 +85,7 @@ def start():
     options = [str(models)]
 
     backend_url = f"{DEFAULT_BACKEND_PROTOCOL}://{host}:{port}"
+    enable_python()
     ctl = Control(options, viasp_backend_url=backend_url)
     for path in paths:
         ctl.load(path)
diff --git a/backend/src/viasp/asp/reify.py b/backend/src/viasp/asp/reify.py
index d05d3285..f4a6820c 100644
--- a/backend/src/viasp/asp/reify.py
+++ b/backend/src/viasp/asp/reify.py
@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Dict, List, Tuple, Iterable, Set, Collection, Any, Union, Sequence
+from typing import Dict, List, Tuple, Iterable, Set, Collection, Any, Union, Sequence, cast
 
 import clingo
 import networkx as nx
@@ -311,11 +311,33 @@ def get_body_aggregate_elements(self, body: Sequence[AST]) -> List[AST]:
         return body_aggregate_elements
 
     def visit_ShowTerm(self, showTerm: AST):
-        new_rule = ast.Rule(
-            showTerm.location, 
-            ast.Literal(showTerm.location, ast.Sign.NoSign, showTerm.term), 
-            showTerm.body)
-        parse_string(place_ast_at_location(new_rule), lambda rule: self.visit(rule))
+        if (hasattr(showTerm, "location") and isinstance(showTerm.location, ast.Location) 
+            and hasattr(showTerm, "term") and isinstance(showTerm.term, AST) 
+            and hasattr(showTerm, "body") and isinstance(showTerm.body, Sequence)
+            and all(isinstance(elem, AST) for elem in showTerm.body)):
+            new_head = ast.Literal(
+                    showTerm.location,
+                    ast.Sign.NoSign,
+                    ast.SymbolicAtom(
+                        showTerm.term
+                        )
+            )
+            self.visit(
+                ast.Rule(
+                showTerm.location, 
+                new_head, 
+                cast(Sequence, showTerm.body))
+            )
+        else:
+            print(f"Plan B for ShowTerm: {showTerm}", flush=True)
+            new_rule = ast.Rule(
+            cast(ast.Location, showTerm.location), 
+            ast.Literal(
+                cast(ast.Location, showTerm.location), 
+                ast.Sign.NoSign, 
+                cast(AST, showTerm.term)), 
+            cast(Sequence, showTerm.body))
+            parse_string(place_ast_at_location(new_rule), lambda rule: self.visit(rule))
 
 
     def visit_Minimize(self, minimize: Minimize):