SC2: Python 3.11 compatibility (#1821)

Co-authored-by: Salzkorn

worlds/_sc2common/bot/sc2process.py

@@ -0,0 +1,275 @@
+import asyncio
+import os
+import os.path
+import shutil
+import signal
+import subprocess
+import sys
+import tempfile
+import time
+from contextlib import suppress
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import aiohttp
+import portpicker
+from worlds._sc2common.bot import logger
+
+from . import paths, wsl
+from .controller import Controller
+from .paths import Paths
+from .versions import VERSIONS
+
+
+class kill_switch:
+    _to_kill: List[Any] = []
+
+    @classmethod
+    def add(cls, value):
+        logger.debug("kill_switch: Add switch")
+        cls._to_kill.append(value)
+
+    @classmethod
+    def kill_all(cls):
+        logger.info(f"kill_switch: Process cleanup for {len(cls._to_kill)} processes")
+        for p in cls._to_kill:
+            # pylint: disable=W0212
+            p._clean(verbose=False)
+
+
+class SC2Process:
+    """
+    A class for handling SCII applications.
+
+    :param host: hostname for the url the SCII application will listen to
+    :param port: the websocket port the SCII application will listen to
+    :param fullscreen: whether to launch the SCII application in fullscreen or not, defaults to False
+    :param resolution: (window width, window height) in pixels, defaults to (1024, 768)
+    :param placement: (x, y) the distances of the SCII app's top left corner from the top left corner of the screen
+                       e.g. (20, 30) is 20 to the right of the screen's left border, and 30 below the top border
+    :param render:
+    :param sc2_version:
+    :param base_build:
+    :param data_hash:
+    """
+
+    def __init__(
+        self,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        fullscreen: bool = False,
+        resolution: Optional[Union[List[int], Tuple[int, int]]] = None,
+        placement: Optional[Union[List[int], Tuple[int, int]]] = None,
+        render: bool = False,
+        sc2_version: str = None,
+        base_build: str = None,
+        data_hash: str = None,
+    ) -> None:
+        assert isinstance(host, str) or host is None
+        assert isinstance(port, int) or port is None
+
+        self._render = render
+        self._arguments: Dict[str, str] = {"-displayMode": str(int(fullscreen))}
+        if not fullscreen:
+            if resolution and len(resolution) == 2:
+                self._arguments["-windowwidth"] = str(resolution[0])
+                self._arguments["-windowheight"] = str(resolution[1])
+            if placement and len(placement) == 2:
+                self._arguments["-windowx"] = str(placement[0])
+                self._arguments["-windowy"] = str(placement[1])
+
+        self._host = host or os.environ.get("SC2CLIENTHOST", "127.0.0.1")
+        self._serverhost = os.environ.get("SC2SERVERHOST", self._host)
+
+        if port is None:
+            self._port = portpicker.pick_unused_port()
+        else:
+            self._port = port
+        self._used_portpicker = bool(port is None)
+        self._tmp_dir = tempfile.mkdtemp(prefix="SC2_")
+        self._process: subprocess = None
+        self._session = None
+        self._ws = None
+        self._sc2_version = sc2_version
+        self._base_build = base_build
+        self._data_hash = data_hash
+
+    async def __aenter__(self) -> Controller:
+        kill_switch.add(self)
+
+        def signal_handler(*_args):
+            # unused arguments: signal handling library expects all signal
+            # callback handlers to accept two positional arguments
+            kill_switch.kill_all()
+
+        signal.signal(signal.SIGINT, signal_handler)
+
+        try:
+            self._process = self._launch()
+            self._ws = await self._connect()
+        except:
+            await self._close_connection()
+            self._clean()
+            raise
+
+        return Controller(self._ws, self)
+
+    async def __aexit__(self, *args):
+        logger.exception("async exit")
+        await self._close_connection()
+        kill_switch.kill_all()
+        signal.signal(signal.SIGINT, signal.SIG_DFL)
+
+    @property
+    def ws_url(self):
+        return f"ws://{self._host}:{self._port}/sc2api"
+
+    @property
+    def versions(self):
+        """Opens the versions.json file which origins from
+        https://github.com/Blizzard/s2client-proto/blob/master/buildinfo/versions.json"""
+        return VERSIONS
+
+    def find_data_hash(self, target_sc2_version: str) -> Optional[str]:
+        """ Returns the data hash from the matching version string. """
+        version: dict
+        for version in self.versions:
+            if version["label"] == target_sc2_version:
+                return version["data-hash"]
+        return None
+
+    def _launch(self):
+        if self._base_build:
+            executable = str(paths.latest_executeble(Paths.BASE / "Versions", self._base_build))
+        else:
+            executable = str(Paths.EXECUTABLE)
+        if self._port is None:
+            self._port = portpicker.pick_unused_port()
+            self._used_portpicker = True
+        args = paths.get_runner_args(Paths.CWD) + [
+            executable,
+            "-listen",
+            self._serverhost,
+            "-port",
+            str(self._port),
+            "-dataDir",
+            str(Paths.BASE),
+            "-tempDir",
+            self._tmp_dir,
+        ]
+        for arg, value in self._arguments.items():
+            args.append(arg)
+            args.append(value)
+        if self._sc2_version:
+
+            def special_match(strg: str):
+                """ Tests if the specified version is in the versions.py dict. """
+                for version in self.versions:
+                    if version["label"] == strg:
+                        return True
+                return False
+
+            valid_version_string = special_match(self._sc2_version)
+            if valid_version_string:
+                self._data_hash = self.find_data_hash(self._sc2_version)
+                assert (
+                    self._data_hash is not None
+                ), f"StarCraft 2 Client version ({self._sc2_version}) was not found inside sc2/versions.py file. Please check your spelling or check the versions.py file."
+
+            else:
+                logger.warning(
+                    f'The submitted version string in sc2.rungame() function call (sc2_version="{self._sc2_version}") was not found in versions.py. Running latest version instead.'
+                )
+
+        if self._data_hash:
+            args.extend(["-dataVersion", self._data_hash])
+
+        if self._render:
+            args.extend(["-eglpath", "libEGL.so"])
+
+        # if logger.getEffectiveLevel() <= logging.DEBUG:
+        args.append("-verbose")
+
+        sc2_cwd = str(Paths.CWD) if Paths.CWD else None
+
+        if paths.PF in {"WSL1", "WSL2"}:
+            return wsl.run(args, sc2_cwd)
+
+        return subprocess.Popen(
+            args,
+            cwd=sc2_cwd,
+            # Suppress Wine error messages
+            stderr=subprocess.DEVNULL
+            # , env=run_config.env
+        )
+
+    async def _connect(self):
+        # How long it waits for SC2 to start (in seconds)
+        for i in range(180):
+            if self._process is None:
+                # The ._clean() was called, clearing the process
+                logger.debug("Process cleanup complete, exit")
+                sys.exit()
+
+            await asyncio.sleep(1)
+            try:
+                self._session = aiohttp.ClientSession()
+                ws = await self._session.ws_connect(self.ws_url, timeout=120)
+                # FIXME fix deprecation warning in for future aiohttp version
+                # ws = await self._session.ws_connect(
+                #     self.ws_url, timeout=aiohttp.client_ws.ClientWSTimeout(ws_close=120)
+                # )
+                logger.debug("Websocket connection ready")
+                return ws
+            except aiohttp.client_exceptions.ClientConnectorError:
+                await self._session.close()
+                if i > 15:
+                    logger.debug("Connection refused (startup not complete (yet))")
+
+        logger.debug("Websocket connection to SC2 process timed out")
+        raise TimeoutError("Websocket")
+
+    async def _close_connection(self):
+        logger.info(f"Closing connection at {self._port}...")
+
+        if self._ws is not None:
+            await self._ws.close()
+
+        if self._session is not None:
+            await self._session.close()
+
+    # pylint: disable=R0912
+    def _clean(self, verbose=True):
+        if verbose:
+            logger.info("Cleaning up...")
+
+        if self._process is not None:
+            if paths.PF in {"WSL1", "WSL2"}:
+                if wsl.kill(self._process):
+                    logger.error("KILLED")
+            elif self._process.poll() is None:
+                for _ in range(3):
+                    self._process.terminate()
+                    time.sleep(0.5)
+                    if not self._process or self._process.poll() is not None:
+                        break
+            else:
+                self._process.kill()
+                self._process.wait()
+                logger.error("KILLED")
+            # Try to kill wineserver on linux
+            if paths.PF in {"Linux", "WineLinux"}:
+                # Command wineserver not detected
+                with suppress(FileNotFoundError):
+                    with subprocess.Popen(["wineserver", "-k"]) as p:
+                        p.wait()
+
+        if os.path.exists(self._tmp_dir):
+            shutil.rmtree(self._tmp_dir)
+
+        self._process = None
+        self._ws = None
+        if self._used_portpicker and self._port is not None:
+            portpicker.return_port(self._port)
+            self._port = None
+        if verbose:
+            logger.info("Cleanup complete")