ta-lib-python/tools/tafungen.py at ab99fdb9fe79cbb8962a442dc3302367812ca2a2 · TA-Lib/ta-lib-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
from talib import abstract, get_functions


def generate_code(filename: str = None):
    file = None
    if filename is not None:
        file = open(filename, "w")
        file.write("from talib import abstract\n\n\n")

    for fun_name in sorted(get_functions()):
        fun = abstract.Function(fun_name)
        print("\n")
        function_code = generate_function_code(fun.__dict__)
        print(function_code)
        if file:
            file.write(function_code)
            file.write("\n\n\n")

    if file:
        file.close()


def get_type_name(i: int):
    return {0: "float",
            1: "str",
            2: "int",
            3: "MA_TYPE"
            }[i]


t = "    "


def generate_function_code(d: dict):
    s_def = f"def {d['_Function__info']['name']}(inputs"

    # list_params = []
    for k, v in d['_Function__info']["parameters"].items():
        # if type(v) == list:
        #     list_params.append(k)
        #     s_def += f", {k}=None"
        # else:
        s_def += f", {k}={v}"

    # list_params2 = []
    for k, v in d['_Function__info']["input_names"].items():
        if type(v) == str:
            s_def += f", {k}='{v}'"
        # elif type(v) == list:
        #     s_def += f", {k}=None"
        #     list_params2.append(k)
        else:
            s_def += f", {k}={v}"

    s_def += "):"

    s_def += f'''
    """
    **{d['_Function__info']['display_name']}**

    Group: {d['_Function__info']['group']}
    '''

    if d['_Function__info']['function_flags'] is not None:
        s_def += f"\n{t}Function flags: {'. '.join(d['_Function__info']['function_flags'])}\n"

    s_def += f"\n{t}:param inputs: input values"

    for k, v in d['_Function__input_names'].items():
        name = d['_Function__input_names'][k]['name']
        default_value = d['_Function__input_names'][k]['price_series']
        if type(default_value) == str:
            default_value = f"'{default_value}'"
        s_def += f"\n{t}:param {k}: {name} (**default**: {default_value})"

    for k, v in d['_Function__opt_inputs'].items():
        definition = d['_Function__opt_inputs'][k]['display_name']
        help_text = d['_Function__opt_inputs'][k]['help']
        default_value = d['_Function__opt_inputs'][k]['default_value']
        type_name = get_type_name(d['_Function__opt_inputs'][k]['type'])
        s_def += f"\n{t}:param {k}: {definition} ({help_text} - **default**: {default_value})"
        s_def += f"\n{t}:type {k}: {type_name}"

    if len(d["_Function__info"]["output_names"]) > 1:
        s_def += f'\n{t}:rtype: Tuple'
        s_def += f'\n{t}:returns: {tuple(d["_Function__info"]["output_names"])}'
    else:
        s_def += f'\n{t}:returns: {d["_Function__info"]["output_names"][0]}'

    s_def += f'\n{t}"""'

    # for lp in list_params:
    #     s_def += f"\n{t}if {lp} is None:"
    #     s_def += f"\n{t}{t}{lp} = {d['_Function__info']['parameters'][lp]}"
    #
    # for lp in list_params2:
    #     s_def += f"\n{t}if {lp} is None:"
    #     s_def += f"\n{t}{t}{lp} = {d['_Function__info']['input_names'][lp]}"

    s_def += f"\n{t}return abstract.Function('{d['_Function__info']['name']}'"

    for k, v in d['_Function__info']["parameters"].items():
        s_def += f", {k}={k}"

    for k, v in d['_Function__info']["input_names"].items():
        s_def += f", {k}={k}"

    s_def += ")(inputs)"

    return s_def


if __name__ == '__main__':
    generate_code("tafun.py")