python - 使用 parsimonious 解析势函数的参数

Question

问题最初在 code review 上提出。通过推荐再次在这里询问。

<小时/>

背景

一个forcefield是用于计算复杂系统势能的函数和参数的集合。我有文本文件，其中包含有关力场参数的数据。文本文件分为许多部分，每个部分都遵循相同的格式:

• 用方括号括起来的节标题
• 在下一行，单词indices:后跟一个整数列表。
• 后面是与该部分关联的 1 行或多行参数

这是一个虚构的示例文件来展示该格式。

############################################
# Comments begin with '#'
############################################

[lj_pairs] # Section 1
    indices:    0 2
#  ID      eps    sigma
    1       2.344   1.234   5
    2       4.423   5.313   5
    3       1.573   6.321   5
    4       1.921   11.93   5

[bonds]
indices:    0 1
    2   4.234e-03   11.2
    6   -0.134545   5.7

目标是解析此类文件并将所有信息存储在 dict 中。

<小时/>

目前，我有以下代码来完成我的任务

""" Force-field data reader """

import re
from dataclasses import dataclass, field
from typing import Dict, Iterable, List, TextIO, Tuple, Union, Any


def ff_reader(fname: Union[str, TextIO]) -> Dict[str, "FFSections"]:
    """ Reads data from a force-field file """

    try:
        if _is_string(fname):
            fh = open(fname, mode="r")
            own = True
        else:
            fh = iter(fname)
    except TypeError:
        raise ValueError("fname must be a string or a file handle")

    # All the possible section headers
    keywords = ("lj_pairs", "bonds")  # etc... Long list of possible sections
                                      # Removed for brevity
    re_sections = re.compile(r"^\[(%s)\]$" % "|".join(keywords))
    ff_data = _strip_comments(fh)
    # Empty dict that'll hold all the data.
    final_ff_data = {key: FFSections() for key in keywords}

    # Get first section header
    for line in ff_data:
        match = re.match(re_sections, line)
        if match:
            section = match.group(1)
            in_section_for_first_time = True
            break
        else:
            raise FFReaderError("A valid section header must be the first line in file")
    else:
        raise FFReaderError("No force-field sections exist")

    # Read the rest of the file
    for line in ff_data:

        match = re.match(re_sections, line)

        # If we've encounted a section header the next line must be an index list.
        if in_section_for_first_time:
            if line.split()[0] != "indices:":
                raise FFReaderError(f"Missing index list for section: {section}")
            idx = _validate_indices(line)
            final_ff_data[section].use_idx = idx
            in_section_for_first_time = False
            in_params_for_first_time = True
            continue

        if match and in_params_for_first_time:
            raise FFReaderError(
                f"Section {section} missing parameters"
                + "Sections must contain atleast one type coefficients"
            )

        if match:  # and not in_section_for_first_time and in_params_for_first_time
            section = match.group(1)
            in_section_for_first_time = True
            continue

        params = _validate_params(line)
        final_ff_data[section].coeffs.update([params])
        in_params_for_first_time = False

    # Close the file if we opened it
    if own:
        fh.close()

    for section in final_ff_data.values():
        # coeff must exist if use_idx does
        if section.use_idx is not None:
            assert section.coeffs

    return final_ff_data

def _strip_comments(
    instream: TextIO, comments: Union[str, Iterable[str], None] = "#"
) -> Iterable[str]:
    """ Strip comments from a text IO stream """

    if comments is not None:
        if isinstance(comments, str):
            comments = [comments]
        comments_re = re.compile("|".join(map(re.escape, comments)))
    else:
        comments_re = ".*"
    try:
        for lines in instream.readlines():
            line = re.split(comments_re, lines, 1)[0].strip()
            if line != "":
                yield line
    except AttributeError:
        raise TypeError("instream must be a `TextIO` stream") from None


@dataclass(eq=False)
class FFSections:
    """
    FFSections(coeffs,use_idx)

    Container for forcefield information
    """

    coeffs: Dict[int, List[float]] = field(default_factory=dict)
    use_idx: List[int] = field(default=None)


class FFReaderError(Exception):
    """ Incorrect or badly formatted force-Field data """

    def __init__(self, message: str, badline: Optional[str] = None) -> None:
        if badline:
            message = f"{message}\nError parsing --> ({badline})"
        super().__init__(message)


def _validate_indices(line: str) -> List[int]:
    """
    Check if given line contains only a whitespace separated
    list of integers
    """
    # split on indices: followed by whitescape
    split = line.split("indices:")[1].split()
    # import ipdb; ipdb.set_trace()
    if not set(s.isdecimal() for s in split) == {True}:
        raise FFReaderError(
            "Indices should be integers and seperated by whitespace", line
        )
    return [int(x) for x in split]


def _validate_params(line: str) -> Tuple[int, List[float]]:
    """
    Check if given line is valid param line, which are
    an integer followed by one or more floats seperated by whitespace
    """
    split = line.split()
    id_ = split[0]
    coeffs = split[1:]
    if not id_.isdecimal():
        raise FFReaderError("Invalid params", line)
    try:
        coeffs = [float(x) for x in coeffs]
    except (TypeError, ValueError):
        raise FFReaderError("Invalid params", line) from None
    return (int(id_), coeffs)

<小时/>

问题

这似乎需要很多代码来完成一个简单的任务。如何使用 parsimonious 或类似的解析库来简化解析此类文件？

Answer 1

正如另一个答案中所述，您可以使用解析库，例如 parsimonious与 NodeVisitor 结合使用类:

from parsimonious.grammar import Grammar
from parsimonious.nodes import NodeVisitor

data = """
############################################
# Comments begin with '#'
############################################

[lj_pairs] # Section 1
    indices:    0 2
    #  ID      eps    sigma
    1       2.344   1.234   5
    2       4.423   5.313   5
    3       1.573   6.321   5
    4       1.921   11.93   5

[bonds]
indices:    0 1
    2   4.234e-03   11.2
    6   -0.134545   5.7
"""

grammar = Grammar(
    r"""
    expr        = (entry / garbage)+
    entry       = section garbage indices (valueline / garbage)*
    section     = lpar word rpar

    indices     = ws? "indices:" values+
    garbage     = ((comment / hs)* newline?)*

    word        = ~"\w+"

    values      = number+
    valueline   = values newline?

    number      = hs? ~"[-.e\d]+" hs?

    lpar        = "["
    rpar        = "]"

    comment     = ~"#.+"
    ws          = ~"\s*"
    hs          = ~"[\t\ ]*"

    newline     = ~"[\r\n]"
    """
)

tree = grammar.parse(data)

class DataVisitor(NodeVisitor):
    def visit_number(self, node, visited_children):
        """ Returns integer and float values. """
        _, value, _ = visited_children
        try:
            number = int(value.text)
        except ValueError:
            number = float(value.text)
        return number

    def visit_section(self, node, visited_children):
        """ Returns the section as text. """
        _, section, _ = visited_children
        return section.text

    def visit_indices(self, node, visited_children):
        """ Returns the index numbers. """
        *_, values = visited_children
        return values[0]

    def visit_valueline(self, node, visited_children):
        """ Returns every value from one line. """
        values, _ = visited_children
        return values

    def visit_entry(self, node, visited_children):
        """ Returns one entry (section, indices, values). """
        section, _, indices, lst = visited_children
        values = [item[0] for item in lst if item[0]]

        return (section, {'indices': indices, 'values': values})

    def visit_expr(self, node, visited_children):
        """ Returns the whole structure as a dict. """
        return dict([item[0] for item in visited_children if item[0]])

    def visit_garbage(self, node, visited_children):
        """ You know what this does. """
        return None

    def generic_visit(self, node, visited_children):
        """ Returns the visited children (if any) or the node itself. """
        return visited_children or node

d = DataVisitor()
result = d.visit(tree)
print(result)

这将产生

{
 'lj_pairs': {'indices': [0, 2], 'values': [[1, 2.344, 1.234, 5], [2, 4.423, 5.313, 5], [3, 1.573, 6.321, 5], [4, 1.921, 11.93, 5]]}, 
 'bonds': {'indices': [0, 1], 'values': [[2, 0.004234, 11.2], [6, -0.134545, 5.7]]}
}

---

说明

您的原始数据文件可以被视为 DSL - 域s特定l语言。因此，我们需要一个语法来描述您的格式允许的样子。这里通常的方法是首先构建小块，例如空格或“单词”。

<小时/>在 parsimonious 我们有几个选项，其中一个是指定正则表达式(以 ~ 开头):

ws          = ~"\s*"

在这里，ws代表\s*这是零个或多个空格。

<小时/>另一种可能性是从字面上形成一个部分，例如

lpar        = "["

<小时/>最后也是最强大的可能性是将这两个较小的部分 组合以形成一个更大的部分，例如

section     = lpar word rpar

翻译为 [word_characters_HERE123]或类似的结构。

<小时/>现在应用正常的交替( / )和量词，例如 * (零矿多，贪心)， + (多一矿石，贪心)和 ? (零或一，贪婪)并且可以放在我们可能想到的每个表达式之后。
<小时/>如果一切正常并且语法适合我们拥有的数据，则所有内容都会被解析为树结构，即所谓的 abstract syntax t ree (AST)。为了真正做某事。对于这个结构很有用(例如，用它制作一个漂亮的 dict )，我们需要将它输入到 NodeVisitor 中。类(class)。这是我们之前形成的语法的附属方法 visit_*会调用每一片适合它的叶子。就是说一个方法visit_section(...)将在每个 section 上被调用叶与其适当的visited_children 。

让我们更清楚地说明这一点。函数

    def visit_section(self, node, visited_children):
        """ Returns the section as text. """
        _, section, _ = visited_children
        return section.text

将调用 section我们语法的一部分( lpar section rpar )，所以叶子 section有这三个 child 。我们对 [ 都不感兴趣也不是 ]但只有部分文本本身，所以我们进行一些解包并返回 section.text 。

我们需要对之前定义的每个节点/叶子执行此操作。默认情况下，第一个定义(在我们的例子中为 expr )和相应的 visit_expr(...)将是 NodeVisitor 的输出class 和所有其他节点都是该节点的子节点(孙子节点、曾孙节点等)。