Added new files

dist/ba_data/python/ba/_assetmanager.py

@@ -0,0 +1,241 @@
+# Released under the MIT License. See LICENSE for details.
+#
+"""Functionality related to managing cloud based assets."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated
+from dataclasses import dataclass, field
+from pathlib import Path
+import threading
+import urllib.request
+import logging
+import weakref
+import time
+import os
+import sys
+
+from efro.dataclassio import (
+    ioprepped,
+    IOAttrs,
+    dataclass_from_json,
+    dataclass_to_json,
+)
+import _ba
+
+if TYPE_CHECKING:
+    from bacommon.assets import AssetPackageFlavor
+
+
+@ioprepped
+@dataclass
+class FileValue:
+    """State for an individual file."""
+
+
+@ioprepped
+@dataclass
+class State:
+    """Holds all persistent state for the asset-manager."""
+
+    files: Annotated[dict[str, FileValue], IOAttrs('files')] = field(
+        default_factory=dict
+    )
+
+
+class AssetManager:
+    """Wrangles all assets."""
+
+    _state: State
+
+    def __init__(self, rootdir: Path) -> None:
+        print('AssetManager()')
+        assert isinstance(rootdir, Path)
+        self.thread_ident = threading.get_ident()
+        self._rootdir = rootdir
+        self._started = False
+        if not self._rootdir.is_dir():
+            raise RuntimeError(f'Provided rootdir does not exist: "{rootdir}"')
+
+        self.load_state()
+
+    def __del__(self) -> None:
+        print('~AssetManager()')
+        if self._started:
+            logging.warning('AssetManager dying in a started state.')
+
+    def launch_gather(
+        self,
+        packages: list[str],
+        flavor: AssetPackageFlavor,
+        account_token: str,
+    ) -> AssetGather:
+        """Spawn an asset-gather operation from this manager."""
+        print(
+            'would gather',
+            packages,
+            'and flavor',
+            flavor,
+            'with token',
+            account_token,
+        )
+        return AssetGather(self)
+
+    def update(self) -> None:
+        """Can be called periodically to perform upkeep."""
+
+    def start(self) -> None:
+        """Tell the manager to start working.
+
+        This will initiate network activity and other processing.
+        """
+        if self._started:
+            logging.warning('AssetManager.start() called on running manager.')
+        self._started = True
+
+    def stop(self) -> None:
+        """Tell the manager to stop working.
+
+        All network activity should be ceased before this function returns.
+        """
+        if not self._started:
+            logging.warning('AssetManager.stop() called on stopped manager.')
+        self._started = False
+        self.save_state()
+
+    @property
+    def rootdir(self) -> Path:
+        """The root directory for this manager."""
+        return self._rootdir
+
+    @property
+    def state_path(self) -> Path:
+        """The path of the state file."""
+        return Path(self._rootdir, 'state')
+
+    def load_state(self) -> None:
+        """Loads state from disk. Resets to default state if unable to."""
+        print('ASSET-MANAGER LOADING STATE')
+        try:
+            state_path = self.state_path
+            if state_path.exists():
+                with open(self.state_path, encoding='utf-8') as infile:
+                    self._state = dataclass_from_json(State, infile.read())
+                    return
+        except Exception:
+            logging.exception('Error loading existing AssetManager state')
+        self._state = State()
+
+    def save_state(self) -> None:
+        """Save state to disk (if possible)."""
+
+        print('ASSET-MANAGER SAVING STATE')
+        try:
+            with open(self.state_path, 'w', encoding='utf-8') as outfile:
+                outfile.write(dataclass_to_json(self._state))
+        except Exception:
+            logging.exception('Error writing AssetManager state')
+
+
+class AssetGather:
+    """Wrangles a gathering of assets."""
+
+    def __init__(self, manager: AssetManager) -> None:
+        assert threading.get_ident() == manager.thread_ident
+        self._manager = weakref.ref(manager)
+        # self._valid = True
+        print('AssetGather()')
+        # url = 'https://files.ballistica.net/bombsquad/promo/BSGamePlay.mov'
+        # url = 'http://www.python.org/ftp/python/2.7.3/Python-2.7.3.tgz'
+        # fetch_url(url,
+        #           filename=Path(manager.rootdir, 'testdl'),
+        #           asset_gather=self)
+        # print('fetch success')
+        thread = threading.Thread(target=self._run)
+        thread.run()
+
+    def _run(self) -> None:
+        """Run the gather in a background thread."""
+        print('hello from gather bg')
+
+        # First, do some sort of.
+
+    # @property
+    # def valid(self) -> bool:
+    #     """Whether this gather is still valid.
+
+    #     A gather becomes in valid if its originating AssetManager dies.
+    #     """
+    #     return True
+
+    def __del__(self) -> None:
+        print('~AssetGather()')
+
+
+def fetch_url(url: str, filename: Path, asset_gather: AssetGather) -> None:
+    """Fetch a given url to a given filename for a given AssetGather."""
+    # pylint: disable=consider-using-with
+    import socket
+
+    # We don't want to keep the provided AssetGather alive, but we want
+    # to abort if it dies.
+    assert isinstance(asset_gather, AssetGather)
+    # weak_gather = weakref.ref(asset_gather)
+
+    # Pass a very short timeout to urllib so we have opportunities
+    # to cancel even with network blockage.
+    req = urllib.request.urlopen(url, context=_ba.app.net.sslcontext, timeout=1)
+    file_size = int(req.headers['Content-Length'])
+    print(f'\nDownloading: {filename} Bytes: {file_size:,}')
+
+    def doit() -> None:
+        time.sleep(1)
+        print('dir', type(req.fp), dir(req.fp))
+        print('WOULD DO IT', flush=True)
+        # req.close()
+        # req.fp.close()
+
+    threading.Thread(target=doit).run()
+
+    with open(filename, 'wb') as outfile:
+        file_size_dl = 0
+
+        block_sz = 1024 * 1024 * 1000
+        time_outs = 0
+        while True:
+            try:
+                data = req.read(block_sz)
+            except ValueError:
+                import traceback
+
+                traceback.print_exc()
+                print('VALUEERROR', flush=True)
+                break
+            except socket.timeout:
+                print('TIMEOUT', flush=True)
+                # File has not had activity in max seconds.
+                if time_outs > 3:
+                    print('\n\n\nsorry -- try back later')
+                    os.unlink(filename)
+                    raise
+                print(
+                    '\nHmmm... little issue... '
+                    'I\'ll wait a couple of seconds'
+                )
+                time.sleep(3)
+                time_outs += 1
+                continue
+
+            # We reached the end of the download!
+            if not data:
+                sys.stdout.write('\rDone!\n\n')
+                sys.stdout.flush()
+                break
+
+            file_size_dl += len(data)
+            outfile.write(data)
+            percent = file_size_dl * 1.0 / file_size
+            status = f'{file_size_dl:20,} Bytes [{percent:.2%}] received'
+            sys.stdout.write('\r' + status)
+            sys.stdout.flush()
+    print('done with', req.fp)