amaranth/nmigen/build/run.py
2020-08-26 22:49:49 +00:00

269 lines
10 KiB
Python

from collections import OrderedDict
from contextlib import contextmanager
from abc import ABCMeta, abstractmethod
import os
import sys
import subprocess
import tempfile
import zipfile
import hashlib
import pathlib
__all__ = ["BuildPlan", "BuildProducts", "LocalBuildProducts", "RemoteSSHBuildProducts"]
class BuildPlan:
def __init__(self, script):
"""A build plan.
Parameters
----------
script : str
The base name (without extension) of the script that will be executed.
"""
self.script = script
self.files = OrderedDict()
def add_file(self, filename, content):
"""
Add ``content``, which can be a :class:`str`` or :class:`bytes`, to the build plan
as ``filename``. The file name can be a relative path with directories separated by
forward slashes (``/``).
"""
assert isinstance(filename, str) and filename not in self.files
self.files[filename] = content
def digest(self, size=64):
"""
Compute a `digest`, a short byte sequence deterministically and uniquely identifying
this build plan.
"""
hasher = hashlib.blake2b(digest_size=size)
for filename in sorted(self.files):
hasher.update(filename.encode("utf-8"))
content = self.files[filename]
if isinstance(content, str):
content = content.encode("utf-8")
hasher.update(content)
hasher.update(self.script.encode("utf-8"))
return hasher.digest()
def archive(self, file):
"""
Archive files from the build plan into ``file``, which can be either a filename, or
a file-like object. The produced archive is deterministic: exact same files will
always produce exact same archive.
"""
with zipfile.ZipFile(file, "w") as archive:
# Write archive members in deterministic order and with deterministic timestamp.
for filename in sorted(self.files):
archive.writestr(zipfile.ZipInfo(filename), self.files[filename])
def execute_local(self, root="build", *, run_script=True):
"""
Execute build plan using the local strategy. Files from the build plan are placed in
the build root directory ``root``, and, if ``run_script`` is ``True``, the script
appropriate for the platform (``{script}.bat`` on Windows, ``{script}.sh`` elsewhere) is
executed in the build root.
Returns :class:`LocalBuildProducts`.
"""
os.makedirs(root, exist_ok=True)
cwd = os.getcwd()
try:
os.chdir(root)
for filename, content in self.files.items():
filename = pathlib.Path(filename)
# Forbid parent directory components completely to avoid the possibility
# of writing outside the build root.
assert ".." not in filename.parts
dirname = os.path.dirname(filename)
if dirname:
os.makedirs(dirname, exist_ok=True)
mode = "wt" if isinstance(content, str) else "wb"
with open(filename, mode) as f:
f.write(content)
if run_script:
if sys.platform.startswith("win32"):
# Without "call", "cmd /c {}.bat" will return 0.
# See https://stackoverflow.com/a/30736987 for a detailed explanation of why.
# Running the script manually from a command prompt is unaffected.
subprocess.check_call(["cmd", "/c", "call {}.bat".format(self.script)])
else:
subprocess.check_call(["sh", "{}.sh".format(self.script)])
return LocalBuildProducts(os.getcwd())
finally:
os.chdir(cwd)
def execute_remote_ssh(self, *, connect_to = {}, root, run_script=True):
"""
Execute build plan using the remote SSH strategy. Files from the build
plan are transferred via SFTP to the directory ``root`` on a remote
server. If ``run_script`` is ``True``, the ``paramiko`` SSH client will
then run ``{script}.sh``. ``root`` can either be an absolute or
relative (to the login directory) path.
``connect_to`` is a dictionary that holds all input arguments to
``paramiko``'s ``SSHClient.connect``
(`documentation <http://docs.paramiko.org/en/stable/api/client.html#paramiko.client.SSHClient.connect>`_).
At a minimum, the ``hostname`` input argument must be supplied in this
dictionary as the remote server.
Returns :class:`RemoteSSHBuildProducts`.
"""
from paramiko import SSHClient
with SSHClient() as client:
client.load_system_host_keys()
client.connect(**connect_to)
with client.open_sftp() as sftp:
def mkdir_exist_ok(path):
try:
sftp.mkdir(str(path))
except IOError as e:
# mkdir fails if directory exists. This is fine in nmigen.build.
# Reraise errors containing e.errno info.
if e.errno:
raise e
def mkdirs(path):
# Iteratively create parent directories of a file by iterating over all
# parents except for the root ("."). Slicing the parents results in
# TypeError, so skip over the root ("."); this also handles files
# already in the root directory.
for parent in reversed(path.parents):
if parent == pathlib.PurePosixPath("."):
continue
else:
mkdir_exist_ok(parent)
mkdir_exist_ok(root)
sftp.chdir(root)
for filename, content in self.files.items():
filename = pathlib.PurePosixPath(filename)
assert ".." not in filename.parts
mkdirs(filename)
mode = "wt" if isinstance(content, str) else "wb"
with sftp.file(str(filename), mode) as f:
# "b/t" modifier ignored in SFTP.
if mode == "wt":
f.write(content.encode("utf-8"))
else:
f.write(content)
if run_script:
transport = client.get_transport()
channel = transport.open_session()
channel.set_combine_stderr(True)
cmd = "if [ -f ~/.profile ]; then . ~/.profile; fi && cd {} && sh {}.sh".format(root, self.script)
channel.exec_command(cmd)
# Show the output from the server while products are built.
buf = channel.recv(1024)
while buf:
print(buf.decode("utf-8"), end="")
buf = channel.recv(1024)
return RemoteSSHBuildProducts(connect_to, root)
def execute(self):
"""
Execute build plan using the default strategy. Use one of the ``execute_*`` methods
explicitly to have more control over the strategy.
"""
return self.execute_local()
class BuildProducts(metaclass=ABCMeta):
@abstractmethod
def get(self, filename, mode="b"):
"""
Extract ``filename`` from build products, and return it as a :class:`bytes` (if ``mode``
is ``"b"``) or a :class:`str` (if ``mode`` is ``"t"``).
"""
assert mode in ("b", "t")
@contextmanager
def extract(self, *filenames):
"""
Extract ``filenames`` from build products, place them in an OS-specific temporary file
location, with the extension preserved, and delete them afterwards. This method is used
as a context manager, e.g.: ::
with products.extract("bitstream.bin", "programmer.cfg") \
as bitstream_filename, config_filename:
subprocess.check_call(["program", "-c", config_filename, bitstream_filename])
"""
files = []
try:
for filename in filenames:
# On Windows, a named temporary file (as created by Python) is not accessible to
# others if it's still open within the Python process, so we close it and delete
# it manually.
file = tempfile.NamedTemporaryFile(
prefix="nmigen_", suffix="_" + os.path.basename(filename),
delete=False)
files.append(file)
file.write(self.get(filename))
file.close()
if len(files) == 0:
return (yield)
elif len(files) == 1:
return (yield files[0].name)
else:
return (yield [file.name for file in files])
finally:
for file in files:
os.unlink(file.name)
class LocalBuildProducts(BuildProducts):
def __init__(self, root):
# We provide no guarantees that files will be available on the local filesystem (i.e. in
# any way other than through `products.get()`) in general, so downstream code must never
# rely on this, even when we happen to use a local build most of the time.
self.__root = root
def get(self, filename, mode="b"):
super().get(filename, mode)
with open(os.path.join(self.__root, filename), "r" + mode) as f:
return f.read()
class RemoteSSHBuildProducts(BuildProducts):
def __init__(self, connect_to, root):
self.__connect_to = connect_to
self.__root = root
def get(self, filename, mode="b"):
super().get(filename, mode)
from paramiko import SSHClient
with SSHClient() as client:
client.load_system_host_keys()
client.connect(**self.__connect_to)
with client.open_sftp() as sftp:
sftp.chdir(self.__root)
with sftp.file(filename, "r" + mode) as f:
# "b/t" modifier ignored in SFTP.
if mode == "t":
return f.read().decode("utf-8")
else:
return f.read()