"""
Design discovery and execution
==============================
This module provides tools for discovering, building, and executing
JITX designs, including communication with the JITX runtime. Normally this
module is not used directly, but rather through the `jitx` command line, or
through the VSCode extension.
.. warning::
The API in this module is still experimental and may change significantly
without notice.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
import asyncio
from collections.abc import Callable, Mapping, Sequence
from logging import getLogger
from typing import Any, overload, override, cast
import jitx.design
from jitx._formatters import Formatter, text_formatter
from jitx.error import InstantiationException, UserCodeException
from jitx.run.runtime import Runtime
from google.protobuf.json_format import MessageToDict
from .._websocket import (
PersistentWebSocketClient,
set_websocket_uri as _set_websocket_uri,
)
from .pyproject import PyProject
logging = getLogger(__name__)
STATUS = "status"
MESSAGE = "message"
ERRORS = "errors"
IMPORTFAILED = "import failed"
INSTANTIATIONFAILED = "instantiation failed"
TRANSLATEFAILED = "translation failed"
DESIGNFAILED = "design failed"
HINTS = "hints"
LOG = "log"
OK = "ok"
[docs]
class DesignFactory:
def __init__(
self,
finder: DesignFinder | None = None,
builder: BaseDesignBuilder | None = None,
*,
formatter: Formatter | None = None,
dump: str | None = None,
):
self.finder = finder or DesignFinder()
self.builder = builder or DesignBuilder()
self.formatter: Formatter = formatter or text_formatter
self.dump = dump
self.queue: dict[str, type[jitx.design.Design] | jitx.design.Design] = {}
self.success = True
[docs]
def by_name(self, name: str):
self.add(self.finder.find_by_name(name))
[docs]
def by_file(self, name: str):
for design in self.finder.find_by_file(name):
self.add(design)
[docs]
def add_all(self):
for design in self.finder.find_all():
if hasattr(design, "__signature__"):
if design.__signature__.parameters:
try:
design.__signature__.bind()
except Exception:
# parameterized design with non-default arguments, skip
continue
self.add(design)
[docs]
def add(
self,
design: type[jitx.design.Design] | jitx.design.Design,
*,
name: str | None = None,
):
if not name:
import jitx._structural
if isinstance(design, jitx._structural.Instantiable):
cls = design._instantiable_()
if isinstance(cls, type) and issubclass(cls, jitx.design.Design):
name = cls.__module__ + "." + cls.__name__
# fake the type to avoid type errors, it's a deferred instantiable of a Design
design = cast(jitx.design.Design, design)
else:
raise ValueError(f"Invalid design: {design}")
elif isinstance(design, jitx.design.Design):
name = design.__class__.__module__ + "." + design.__class__.__name__
elif isinstance(design, type) and issubclass(design, jitx.design.Design):
name = design.__module__ + "." + design.__name__
else:
raise ValueError(f"Invalid design: {design}")
if name in self.queue:
raise ValueError(f"Duplicate design added to build queue: {name}")
self.queue[name] = design
[docs]
def build(self):
aggregate = {}
for name, design in self.queue.items():
result = self.builder.build(
design, name=name, dump=self.dump, formatter=self.formatter
)
aggregate[name] = result
if ERRORS in result or result.get(STATUS) != OK:
self.success = False
if self.finder.exceptions:
self.success = False
aggregate[ERRORS] = {
IMPORTFAILED: {
name: repr(e) for name, e in self.finder.exceptions.items()
}
}
self.formatter(aggregate)
[docs]
def list(self):
result = {}
df = self.finder
result["designs"] = [d.__module__ + "." + d.__name__ for d in df.find_all()]
if df.exceptions:
if ERRORS not in result:
result[ERRORS] = {}
if IMPORTFAILED not in result[ERRORS]:
result[ERRORS][IMPORTFAILED] = {}
result[ERRORS][IMPORTFAILED].update(
(name, repr(e)) for name, e in df.exceptions.items()
)
self.formatter(result)
[docs]
class DesignFinder:
def __init__(self, roots: str | Sequence[str] | None = None):
if isinstance(roots, str):
roots = (roots,)
self.roots = roots or (".",)
self.exceptions: dict[str, Exception] = {}
[docs]
def find_all(self):
import os
for root in self.roots:
project = PyProject(root)
tool = project.jitxtool
for dirpath, dirnames, filenames in os.walk(root):
path = os.path.relpath(dirpath, root)
if path in tool.exclude:
dirnames[:] = []
else:
dirnames[:] = [
d
for d in dirnames
if not d.startswith("_")
and not d.startswith(".")
and not any(d == ex for ex in tool.exclude)
]
for filename in filenames:
if filename.startswith("_") or filename.startswith("."):
continue
if any(filename == ex for ex in tool.exclude):
continue
if not filename.endswith(".py"):
continue
yield from self.find_by_file(os.path.join(path, filename))
# NOTE walk_packages struggles with namespace packages, so we have to
# resort to look for python files for now. This doesn't support things
# like eggs, which is probably fine anyway.
# import pkgutil
# for mi in pkgutil.walk_packages(self.roots):
# try:
# yield from self.find_by_module(mi.name)
# except Exception as e:
# self.exceptions[mi.name] = e
[docs]
def find_by_name(self, name: str):
import importlib
ns = name.rsplit(".", 1)
if len(ns) != 2:
raise ValueError(f"Invalid design name: {name}")
modulename, classname = ns
m = importlib.import_module(modulename)
design = getattr(m, classname, None)
if design is None:
raise ValueError(f"{classname} not found in {modulename}")
if not issubclass(design, jitx.design.Design):
raise ValueError(f"{classname} in {modulename} is not a Design")
return design
[docs]
def find_by_module(self, name: str):
import importlib
import jitx.sample
try:
m = importlib.import_module(name)
except Exception as e:
self.exceptions[name] = e
return
for elem in dir(m):
field = getattr(m, elem, None)
if (
isinstance(field, type)
and issubclass(field, jitx.design.Design)
and field not in (jitx.design.Design, jitx.sample.SampleDesign)
and field.__module__ == m.__name__
):
yield field
[docs]
def find_by_file(self, path: str):
import os.path
import importlib
path, filename = os.path.split(os.path.normpath(path))
module, _ = os.path.splitext(filename)
steps = []
while path:
rem, last = os.path.split(path)
if rem == path:
break
path = rem
steps.append(last)
steps.reverse()
steps.append(module)
# attempt to find the longest matching module path to avoid accidentally
# importing a "shadowed" module
candidate = module
for i in range(len(steps)):
if steps[i] == "src":
# do not accept "src" as top-level module, this is common in the
# so called "src-layout" pattern, and it could technically be a
# valid top level package, it's exceedingly unlikely.
continue
candidate = ".".join(steps[i:])
try:
importlib.import_module(candidate)
break
except ModuleNotFoundError as e:
if e.name != candidate and e.name != steps[i]:
self.exceptions[candidate] = e
# only continue trying if the error is about the candidate,
# otherwise it's probably an error _inside_ the module.
return
except Exception as e:
# some other error happening on import here, so we're probably
# in the right spot.
self.exceptions[candidate] = e
return
yield from self.find_by_module(candidate)
[docs]
class BaseDesignBuilder(ABC):
@abstractmethod
def __init__(self):
pass
[docs]
def build(
self,
design: type[jitx.design.Design] | jitx.design.Design,
*,
name: str | None = None,
dump: str | None = None,
formatter: Formatter,
) -> Mapping[str, Any]:
result: dict[str, Any] = {"design": name}
try:
dumper: None | Callable = None
if dump:
def dumpfn(x):
with open(dump, "w") as f:
formatter(MessageToDict(x), file=f)
dumper = dumpfn
result.update(
asyncio.run(self._submit_design(design, name=name, dump=dumper))
)
except UserCodeException as e:
result[ERRORS] = {TRANSLATEFAILED: [str(e)]}
if e.hint:
result[HINTS] = [e.hint]
return result
except Exception as e:
logging.exception(f"Unable to submit {design}")
errors: list[BaseException] = [e]
while e.__cause__ is not None:
e = e.__cause__
errors.insert(0, e)
if (tb := e.__traceback__) is not None:
tb = tb.tb_next
def formatexc(e: BaseException):
if isinstance(e, UserCodeException | InstantiationException):
return str(e)
loc = ""
if tb := e.__traceback__:
while tb.tb_next:
tb = tb.tb_next
loc = f" at {tb.tb_frame.f_code.co_filename}:{tb.tb_lineno}"
pass
return str(e) + loc
result[ERRORS] = {DESIGNFAILED: [formatexc(e) for e in errors]}
return result
return result
def _instantiate_design(
self,
design: type[jitx.design.Design] | jitx.design.Design,
):
import jitx._structural
from jitx._instantiation import instantiation
import jitx._translate.design
import gc
with instantiation.activate():
if isinstance(design, jitx._structural.Instantiable):
cls = design._instantiable_()
if not (isinstance(cls, type) and issubclass(cls, jitx.design.Design)):
raise ValueError(f"Not an instantiable design: {cls}")
with instantiation.frame():
instantiated = design._instantiate_({})
elif isinstance(design, type):
if not issubclass(design, jitx.design.Design):
raise ValueError(f"Not a Design subclass: {cls}")
try:
if hasattr(design, "__signature__"):
design.__signature__.bind()
except TypeError:
raise ValueError(
"Design is parameterized but no parameters were provided"
) from None
with instantiation.frame():
instantiated = design()
else:
instantiated = design
assert isinstance(instantiated, jitx.design.Design), (
f"Got non-Design object {jitx._structural.Proxy.type(instantiated)}"
)
# try to force detection of lost elements.
gc.collect()
return instantiated
async def _submit_design(
self,
design: jitx.design.Design | type[jitx.design.Design],
name: str | None,
*,
dump: Callable[[Any], None] | None = None,
) -> Mapping[str, Any]:
raise NotImplementedError
[docs]
class DryRunBuilder(BaseDesignBuilder):
def __init__(self):
super().__init__()
async def _submit_design(
self,
design: jitx.design.Design | type[jitx.design.Design],
name: str | None,
*,
dump: Callable[[Any], None] | None = None,
):
from jitx.substrate import SubstrateContext
from jitx.design import DesignContext
import jitx._translate.design
from jitx._instantiation import instantiation
instantiated = self._instantiate_design(design)
with instantiation.activate():
with DesignContext(instantiated):
with SubstrateContext(instantiated.substrate):
packaged = jitx._translate.design.package_design(instantiated)
if dump:
dump(packaged)
return {STATUS: OK}
[docs]
class DesignBuilder(BaseDesignBuilder):
@overload
def __init__(self, *, spec: str | None = None): ...
@overload
def __init__(self, *, uri: str): ...
@overload
def __init__(self, *, port: int, host: str = "localhost"): ...
__client_future: asyncio.Task[PersistentWebSocketClient] | None = None
def __init__(
self,
*,
spec: str | None = None,
uri: str | None = None,
port: int | None = None,
host: str | None = None,
):
super().__init__()
def lazy_setup():
if uri is not None:
_set_websocket_uri(uri=uri)
elif port is not None:
_set_websocket_uri(uri=f"ws://{host or 'localhost'}:{port}")
elif spec is not None:
_set_websocket_uri(file=spec)
else:
_set_websocket_uri()
self.__setup = lambda: None
self.__setup = lazy_setup
async def __client(self):
if self.__client_future is None:
self.__client_future = asyncio.create_task(
PersistentWebSocketClient.create()
)
return await self.__client_future
[docs]
@override
def build(
self,
design: type[jitx.design.Design] | jitx.design.Design,
*,
name: str | None = None,
dump: str | None = None,
formatter: Formatter,
) -> Mapping[str, Any]:
self.__setup()
return super().build(design, name=name, formatter=formatter, dump=dump)
async def _submit_design(
self,
design: type[jitx.design.Design] | jitx.design.Design,
name: str | None,
*,
dump: Callable[[Any], None] | None = None,
):
client = await self.__client()
await client.establish()
async with Runtime(client=client) as r:
if dump:
r_package_design = r._package_design
def package_design(
constructed: jitx.design.Design, name: str | None = None
):
packaged, rmap = r_package_design(constructed, name)
dump(packaged)
return packaged, rmap
r._package_design = package_design
await r.submit(design, name=name)
return {STATUS: OK}
