nmigen/build/run.py

   1 from collections import OrderedDict
   2 from contextlib import contextmanager
   3 from abc import ABCMeta, abstractmethod
   4 import os
   5 import sys
   6 import subprocess
   7 import tempfile
   8 import zipfile
   9 import hashlib
  10 import pathlib
  11
  12
  13 __all__ = ["BuildPlan", "BuildProducts", "LocalBuildProducts", "RemoteSSHBuildProducts"]
  14
  15
  16
  17 class BuildPlan:
  18     def __init__(self, script):
  19         """A build plan.
  20
  21         Parameters
  22         ----------
  23         script : str
  24             The base name (without extension) of the script that will be executed.
  25         """
  26         self.script = script
  27         self.files  = OrderedDict()
  28
  29     def add_file(self, filename, content):
  30         """
  31         Add ``content``, which can be a :class:`str`` or :class:`bytes`, to the build plan
  32         as ``filename``. The file name can be a relative path with directories separated by
  33         forward slashes (``/``).
  34         """
  35         assert isinstance(filename, str) and filename not in self.files
  36         self.files[filename] = content
  37
  38     def digest(self, size=64):
  39         """
  40         Compute a `digest`, a short byte sequence deterministically and uniquely identifying
  41         this build plan.
  42         """
  43         hasher = hashlib.blake2b(digest_size=size)
  44         for filename in sorted(self.files):
  45             hasher.update(filename.encode("utf-8"))
  46             content = self.files[filename]
  47             if isinstance(content, str):
  48                 content = content.encode("utf-8")
  49             hasher.update(content)
  50         hasher.update(self.script.encode("utf-8"))
  51         return hasher.digest()
  52
  53     def archive(self, file):
  54         """
  55         Archive files from the build plan into ``file``, which can be either a filename, or
  56         a file-like object. The produced archive is deterministic: exact same files will
  57         always produce exact same archive.
  58         """
  59         with zipfile.ZipFile(file, "w") as archive:
  60             # Write archive members in deterministic order and with deterministic timestamp.
  61             for filename in sorted(self.files):
  62                 archive.writestr(zipfile.ZipInfo(filename), self.files[filename])
  63
  64     def execute_local(self, root="build", *, run_script=True):
  65         """
  66         Execute build plan using the local strategy. Files from the build plan are placed in
  67         the build root directory ``root``, and, if ``run_script`` is ``True``, the script
  68         appropriate for the platform (``{script}.bat`` on Windows, ``{script}.sh`` elsewhere) is
  69         executed in the build root.
  70
  71         Returns :class:`LocalBuildProducts`.
  72         """
  73         os.makedirs(root, exist_ok=True)
  74         cwd = os.getcwd()
  75         try:
  76             os.chdir(root)
  77
  78             for filename, content in self.files.items():
  79                 filename = pathlib.Path(filename)
  80                 # Forbid parent directory components completely to avoid the possibility
  81                 # of writing outside the build root.
  82                 assert ".." not in filename.parts
  83                 dirname = os.path.dirname(filename)
  84                 if dirname:
  85                     os.makedirs(dirname, exist_ok=True)
  86
  87                 mode = "wt" if isinstance(content, str) else "wb"
  88                 with open(filename, mode) as f:
  89                     f.write(content)
  90
  91             if run_script:
  92                 if sys.platform.startswith("win32"):
  93                     # Without "call", "cmd /c {}.bat" will return 0.
  94                     # See https://stackoverflow.com/a/30736987 for a detailed explanation of why.
  95                     # Running the script manually from a command prompt is unaffected.
  96                     subprocess.check_call(["cmd", "/c", "call {}.bat".format(self.script)])
  97                 else:
  98                     subprocess.check_call(["sh", "{}.sh".format(self.script)])
  99
 100             return LocalBuildProducts(os.getcwd())
 101
 102         finally:
 103             os.chdir(cwd)
 104
 105     def execute_remote_ssh(self, *, connect_to = {}, root, run_script=True):
 106         """
 107         Execute build plan using the remote SSH strategy. Files from the build
 108         plan are transferred via SFTP to the directory ``root`` on a  remote
 109         server. If ``run_script`` is ``True``, the ``paramiko`` SSH client will
 110         then run ``{script}.sh``. ``root`` can either be an absolute or
 111         relative (to the login directory) path.
 112
 113         ``connect_to`` is a dictionary that holds all input arguments to
 114         ``paramiko``'s ``SSHClient.connect``
 115         (`documentation <http://docs.paramiko.org/en/stable/api/client.html#paramiko.client.SSHClient.connect>`_).
 116         At a minimum, the ``hostname`` input argument must be supplied in this
 117         dictionary as the remote server.
 118
 119         Returns :class:`RemoteSSHBuildProducts`.
 120         """
 121         from paramiko import SSHClient
 122
 123         with SSHClient() as client:
 124             client.load_system_host_keys()
 125             client.connect(**connect_to)
 126
 127             with client.open_sftp() as sftp:
 128                 def mkdir_exist_ok(path):
 129                     try:
 130                         sftp.mkdir(str(path))
 131                     except IOError as e:
 132                         # mkdir fails if directory exists. This is fine in nmigen.build.
 133                         # Reraise errors containing e.errno info.
 134                         if e.errno:
 135                             raise e
 136
 137                 def mkdirs(path):
 138                     # Iteratively create parent directories of a file by iterating over all
 139                     # parents except for the root ("."). Slicing the parents results in
 140                     # TypeError, so skip over the root ("."); this also handles files
 141                     # already in the root directory.
 142                     for parent in reversed(path.parents):
 143                         if parent == pathlib.PurePosixPath("."):
 144                             continue
 145                         else:
 146                             mkdir_exist_ok(parent)
 147
 148                 mkdir_exist_ok(root)
 149
 150                 sftp.chdir(root)
 151                 for filename, content in self.files.items():
 152                     filename = pathlib.PurePosixPath(filename)
 153                     assert ".." not in filename.parts
 154
 155                     mkdirs(filename)
 156
 157                     mode = "wt" if isinstance(content, str) else "wb"
 158                     with sftp.file(str(filename), mode) as f:
 159                         # "b/t" modifier ignored in SFTP.
 160                         if mode == "wt":
 161                             f.write(content.encode("utf-8"))
 162                         else:
 163                             f.write(content)
 164
 165             if run_script:
 166                 transport = client.get_transport()
 167                 channel = transport.open_session()
 168                 channel.set_combine_stderr(True)
 169
 170                 cmd = "if [ -f ~/.profile ]; then . ~/.profile; fi && cd {} && sh {}.sh".format(root, self.script)
 171                 channel.exec_command(cmd)
 172
 173                 # Show the output from the server while products are built.
 174                 buf = channel.recv(1024)
 175                 while buf:
 176                     print(buf.decode("utf-8"), end="")
 177                     buf = channel.recv(1024)
 178
 179         return RemoteSSHBuildProducts(connect_to, root)
 180
 181     def execute(self):
 182         """
 183         Execute build plan using the default strategy. Use one of the ``execute_*`` methods
 184         explicitly to have more control over the strategy.
 185         """
 186         return self.execute_local()
 187
 188
 189 class BuildProducts(metaclass=ABCMeta):
 190     @abstractmethod
 191     def get(self, filename, mode="b"):
 192         """
 193         Extract ``filename`` from build products, and return it as a :class:`bytes` (if ``mode``
 194         is ``"b"``) or a :class:`str` (if ``mode`` is ``"t"``).
 195         """
 196         assert mode in ("b", "t")
 197
 198     @contextmanager
 199     def extract(self, *filenames):
 200         """
 201         Extract ``filenames`` from build products, place them in an OS-specific temporary file
 202         location, with the extension preserved, and delete them afterwards. This method is used
 203         as a context manager, e.g.: ::
 204
 205             with products.extract("bitstream.bin", "programmer.cfg") \
 206                     as bitstream_filename, config_filename:
 207                 subprocess.check_call(["program", "-c", config_filename, bitstream_filename])
 208         """
 209         files = []
 210         try:
 211             for filename in filenames:
 212                 # On Windows, a named temporary file (as created by Python) is not accessible to
 213                 # others if it's still open within the Python process, so we close it and delete
 214                 # it manually.
 215                 file = tempfile.NamedTemporaryFile(
 216                     prefix="nmigen_", suffix="_" + os.path.basename(filename),
 217                     delete=False)
 218                 files.append(file)
 219                 file.write(self.get(filename))
 220                 file.close()
 221
 222             if len(files) == 0:
 223                 return (yield)
 224             elif len(files) == 1:
 225                 return (yield files[0].name)
 226             else:
 227                 return (yield [file.name for file in files])
 228         finally:
 229             for file in files:
 230                 os.unlink(file.name)
 231
 232
 233 class LocalBuildProducts(BuildProducts):
 234     def __init__(self, root):
 235         # We provide no guarantees that files will be available on the local filesystem (i.e. in
 236         # any way other than through `products.get()`) in general, so downstream code must never
 237         # rely on this, even when we happen to use a local build most of the time.
 238         self.__root = root
 239
 240     def get(self, filename, mode="b"):
 241         super().get(filename, mode)
 242         with open(os.path.join(self.__root, filename), "r" + mode) as f:
 243             return f.read()
 244
 245
 246 class RemoteSSHBuildProducts(BuildProducts):
 247     def __init__(self, connect_to, root):
 248         self.__connect_to = connect_to
 249         self.__root = root
 250
 251     def get(self, filename, mode="b"):
 252         super().get(filename, mode)
 253
 254         from paramiko import SSHClient
 255
 256         with SSHClient() as client:
 257             client.load_system_host_keys()
 258             client.connect(**self.__connect_to)
 259
 260             with client.open_sftp() as sftp:
 261                 sftp.chdir(self.__root)
 262
 263                 with sftp.file(filename, "r" + mode) as f:
 264                     # "b/t" modifier ignored in SFTP.
 265                     if mode == "t":
 266                         return f.read().decode("utf-8")
 267                     else:
 268                         return f.read()