pytorch

1
import re
2
from typing import Mapping, Match, Optional, Sequence
3

4
# match $identifier or ${identifier} and replace with value in env
5
# If this identifier is at the beginning of whitespace on a line
6
# and its value is a list then it is treated as
7
# block substitution by indenting to that depth and putting each element
8
# of the list on its own line
9
# if the identifier is on a line starting with non-whitespace and a list
10
# then it is comma separated ${,foo} will insert a comma before the list
11
# if this list is not empty and ${foo,} will insert one after.
12

13

14
class CodeTemplate:
15
    substitution_str = r"(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})"
16
    substitution = re.compile(substitution_str, re.MULTILINE)
17

18
    pattern: str
19
    filename: str
20

21
    @staticmethod
22
    def from_file(filename: str) -> "CodeTemplate":
23
        with open(filename) as f:
24
            return CodeTemplate(f.read(), filename)
25

26
    def __init__(self, pattern: str, filename: str = "") -> None:
27
        self.pattern = pattern
28
        self.filename = filename
29

30
    def substitute(
31
        self, env: Optional[Mapping[str, object]] = None, **kwargs: object
32
    ) -> str:
33
        if env is None:
34
            env = {}
35

36
        def lookup(v: str) -> object:
37
            assert env is not None
38
            return kwargs[v] if v in kwargs else env[v]
39

40
        def indent_lines(indent: str, v: Sequence[object]) -> str:
41
            return "".join(
42
                [indent + l + "\n" for e in v for l in str(e).splitlines()]
43
            ).rstrip()
44

45
        def replace(match: Match[str]) -> str:
46
            indent = match.group(1)
47
            key = match.group(2)
48
            comma_before = ""
49
            comma_after = ""
50
            if key[0] == "{":
51
                key = key[1:-1]
52
                if key[0] == ",":
53
                    comma_before = ", "
54
                    key = key[1:]
55
                if key[-1] == ",":
56
                    comma_after = ", "
57
                    key = key[:-1]
58
            v = lookup(key)
59
            if indent is not None:
60
                if not isinstance(v, list):
61
                    v = [v]
62
                return indent_lines(indent, v)
63
            elif isinstance(v, list):
64
                middle = ", ".join([str(x) for x in v])
65
                if len(v) == 0:
66
                    return middle
67
                return comma_before + middle + comma_after
68
            else:
69
                return str(v)
70

71
        return self.substitution.sub(replace, self.pattern)
72

73

74
if __name__ == "__main__":
75
    c = CodeTemplate(
76
        """\
77
    int foo($args) {
78

79
        $bar
80
            $bar
81
        $a+$b
82
    }
83
    int commatest(int a${,stuff})
84
    int notest(int a${,empty,})
85
    """
86
    )
87
    print(
88
        c.substitute(
89
            args=["hi", 8],
90
            bar=["what", 7],
91
            a=3,
92
            b=4,
93
            stuff=["things...", "others"],
94
            empty=[],
95
        )
96
    )
97