From 6b08d8c8421da3ea2365a6bb55b5848e324f8a30 Mon Sep 17 00:00:00 2001
From: Vadim Liventsev <dev@vadim.me>
Date: Tue, 3 Oct 2023 17:39:05 +0200
Subject: [PATCH] Added pyte to handle terminal schenanigans

---
 .github/workflows/build.yml |  1 +
 programlib/__init__.py      |  2 ++
 programlib/program.py       |  6 +++---
 programlib/terminal.py      | 18 ++++++++++++++++++
 pyproject.toml              |  1 +
 5 files changed, 25 insertions(+), 3 deletions(-)
 create mode 100644 programlib/terminal.py

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 38b9488..e47336a 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -8,4 +8,5 @@ jobs:
       - name: Build and publish to pypi
         uses: JRubics/poetry-publish@master
         with:
+          python_version: "3.10.8"
           pypi_token: ${{ secrets.PYPI_TOKEN }}
\ No newline at end of file
diff --git a/programlib/__init__.py b/programlib/__init__.py
index d66e6b0..6cc1947 100644
--- a/programlib/__init__.py
+++ b/programlib/__init__.py
@@ -1,2 +1,4 @@
 from .language import Language, languages, language_
+from .agent import Agent
+from .terminal import Terminal
 from .program import Program
\ No newline at end of file
diff --git a/programlib/program.py b/programlib/program.py
index 6009b07..be5379b 100644
--- a/programlib/program.py
+++ b/programlib/program.py
@@ -4,9 +4,8 @@
 from uuid import uuid4
 from pathlib import Path
 from itertools import zip_longest
-from programlib.agent import Agent
 
-from programlib import language_
+from programlib import Agent, language_, Terminal
 
 def correctness(expected_outputs, outputs):
     assert expected_outputs, 'expected_outputs is empty. Cannot test program correctness.'
@@ -29,6 +28,7 @@ class Program():
     def __init__(self, source=None, name=None, language='C++', 
                        workdir=Path(__file__).parent / 'programs'):
         self.language = language_(language)
+        self.term = Terminal()
 
         self.workdir = workdir
         self.name = name or str(uuid4())
@@ -62,7 +62,7 @@ def run(self, input_lines=[], force=True):
 
         self.stdout, self.exitstatus = self.language.run(self.workdir, self.name, input_lines)
         assert force or not self.exitstatus, f'Exit status {self.exitstatus}'
-        return self.stdout.splitlines()
+        return repr(self.stdout).splitlines()
     
     def spawn(self, delimiter='\n'):
         """
diff --git a/programlib/terminal.py b/programlib/terminal.py
new file mode 100644
index 0000000..9f63583
--- /dev/null
+++ b/programlib/terminal.py
@@ -0,0 +1,18 @@
+import pyte
+
+class Terminal:
+    def __init__(self, rows=24, cols=80):
+        self.screen = pyte.Screen(cols, rows)
+        self.stream = pyte.Stream(self.screen)
+
+    def emulate(self, raw_output, clean=True):
+        self.stream.feed(raw_output)
+
+        lines = self.screen.display
+        self.screen.reset()
+
+        if clean:
+            lines = (line.rstrip() for line in lines)
+            lines = (line for line in lines if line)
+
+        return lines
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 61d55c4..c5d07fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,7 @@ python = "^3.8"
 pexpect = "^4.8.0"
 gym = "^0.26.2"
 numpy = "^1.24.2"
+pyte = "^0.8.0"
 
 [build-system]
 requires = ["poetry-core"]