Commit c6340958 authored by Martin Larralde's avatar Martin Larralde
Browse files

Update `setup.py` to generate `array` look-up tables instead of using `dict` tables

parent a530e3ef
......@@ -4,6 +4,7 @@
### Added manually
# Build files
peptides/data/tables.py
peptides/data/lut.py
### C ###
# Prerequisites
......
......@@ -5,4 +5,5 @@ include pyproject.toml
recursive-include peptides/data *.csv *.py
include peptides/py.typed
exclude peptides/data/tables.py
exclude peptides/data/lut.py
prune tests
......@@ -75,9 +75,30 @@ class codegen(setuptools.Command):
with open(filename, "r") as f:
return {row[0]:float(row[1]) for row in csv.reader(f)}
def _generate_module(self, tables):
def _load_code_1(self):
filename = os.path.realpath(os.path.join(__file__, "..", "peptides", "__init__.py"))
with open(filename) as f:
contents = f.read()
peptides_module = ast.parse(contents, filename)
peptide_class = next(
expr
for expr in peptides_module.body
if isinstance(expr, ast.ClassDef)
and expr.name == "Peptide"
)
code1 = next(
expr
for expr in peptide_class.body
if isinstance(expr, ast.Assign)
and expr.targets[0].id == "_CODE1"
)
return "".join(x.value for x in code1.value.elts)
def _generate_tables_module(self, tables):
n = sum(map(len, tables.values()))
self.announce(f"building Python AST from {n!r} data tables", level=2)
self.announce(f"building Python AST for {n!r} data tables", level=2)
body = []
for name, table in tables.items():
stub = name.upper().replace(".", "_")
......@@ -97,10 +118,48 @@ class codegen(setuptools.Command):
body.append(assign_node)
return ast.Module(body=body)
def _generate_lut_module(self, tables, code1):
n = sum(map(len, tables.values()))
self.announce(f"building Python AST for {n!r} look-up tables", level=2)
body = [
ast.ImportFrom(module='array', names=[ast.alias(name='array')], level=0)
]
for name, table in tables.items():
stub = name.upper().replace(".", "_")
assign_node = ast.Assign(
targets=[ast.Name(id=stub, ctx=ast.Store())],
value=ast.Dict(
keys=list(map(ast.Constant, table.keys())),
values=[
ast.Call(
func=ast.Name(id='array', ctx=ast.Load()),
args=[
ast.Constant(value='d'),
ast.List(
elts=[
ast.Constant(value=subtable.get(aa, 0.0))
for aa in code1
]
)
],
keywords=[],
)
#
# ast.Dict(
# keys=list(map(ast.Constant, subtable.keys())),
# values=list(map(ast.Constant, subtable.values())),
# )
for subtable in table.values()
]
)
)
body.append(assign_node)
return ast.Module(body=body)
def _write_module(self, module, filename):
self.announce(f"writing Python source to {filename!r}", level=2)
with open(filename, "w") as f:
f.write("# this file was automatically generated")
f.write("# this file was automatically generated\n")
f.write("# by the `python setup.py codegen` command\n")
f.write("# DO NOT EDIT MANUALLY!\n")
f.write(astor.to_source(module))
......@@ -110,18 +169,30 @@ class codegen(setuptools.Command):
# check astor is available
if isinstance(astor, ImportError):
raise RuntimeError("`astor` package is required for the `codegen` command") from astor
# load 1-letter code
code1 = self._load_code_1()
# load data tables from R-formatted data file
tables = self._load_tables()
# generate a Python file containing the constants
module_node = self._generate_module(tables)
tables_module = self._generate_tables_module(tables)
lut_module = self._generate_lut_module(tables, code1)
# write the data sources as Python
output_file = os.path.join(self.build_lib, "peptides", "data", "tables.py")
self.mkpath(os.path.dirname(output_file))
self._write_module(module_node, output_file)
table_file = os.path.join(self.build_lib, "peptides", "data", "tables.py")
self.mkpath(os.path.dirname(table_file))
self._write_module(tables_module, table_file)
# write the data sources as Python
lut_file = os.path.join(self.build_lib, "peptides", "data", "lut.py")
self.mkpath(os.path.dirname(lut_file))
self._write_module(lut_module, lut_file)
# copy if inplace
if self.inplace:
library_file = os.path.join("peptides", "data", "tables.py")
self.copy_file(output_file, library_file)
library_table_file = os.path.join("peptides", "data", "tables.py")
self.copy_file(table_file, library_table_file)
library_lut_file = os.path.join("peptides", "data", "lut.py")
self.copy_file(lut_file, library_lut_file)
class build_py(_build_py):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment