neuralforecast/scripts/cli.py at 3410ce8dc8ed5fde2b309faee40064ff76481006 · Nixtla/neuralforecast · 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
# AUTOGENERATED! DO NOT EDIT! File to edit: cli.ipynb.

# %% auto 0
__all__ = ['tst_flags', 'to_skip', 'mapper', 'print_execs', 'print_hide', 'other_tests', 'get_markdown', 'get_code',
           'extract_dir', 'no_dir_and_dir', 'get_all_tests2', 'get_all_tests', 'print_dir_in_nb']

# %% cli.ipynb 1
from functools import partial

from execnb.nbio import read_nb
from fastcore.script import call_parse
from fastcore.xtras import Path, globtastic
from nbdev import nbdev_export
from nbdev.export import ExportModuleProc, nb_export
from nbdev.maker import ModuleMaker
from nbdev.processors import NBProcessor


# %% cli.ipynb 4
tst_flags = 'datasets distributed matplotlib polars pyarrow scipy'.split()
to_skip = [
    'showdoc',
    'load_ext',
    'from nbdev'
]


def print_execs(cell):
    if 'exec' in cell.source: print(cell.source)

def print_hide(cell):
    if 'hide' in cell.directives_: print(cell.source)

def other_tests(cell):
    if len(cell.directives_) == 0:
        print(cell.source)

def get_markdown(cell):
    if cell.cell_type == "markdown":
        print(cell.source)

def get_code(cell):
    if cell.cell_type == "code":
        print(cell.source)

def extract_dir(cell, dirs):
    for directive in dirs.split(','):
        if directive in cell.directives_:
            print(cell.source)

def no_dir_and_dir(cell, dir):
    if len(cell.directives_) == 0:
        print(cell.source)

    if dir in cell.directives_:
        print(cell.source)

def get_all_tests2(cell):
    if cell.cell_type == "code":

        if len(cell.directives_) == 0:
            print(cell.source)


        elif any(x in tst_flags + ['hide'] for x in cell.directives_):
            # if not (x in cell.source for x in to_skip):
            print(cell.source)

def get_all_tests(cell):
    if len(cell.directives_) == 0:
        print(cell.source)

    if any(x in tst_flags + ["hide"] for x in cell.directives_):
        print(cell.source)


# %% cli.ipynb 8
mapper = {
    'print_execs': print_execs,
    'print_hide': print_hide,
    'other_tests': other_tests,
    'get_markdown': get_markdown,
    'extract_dir': extract_dir,
    'no_dir_and_dir': no_dir_and_dir,
    'get_all_tests':get_all_tests,
    'get_all_tests2':get_all_tests2
}

# %% cli.ipynb 9
@call_parse
def print_dir_in_nb(nb_path:str,
                    dir:str=None,
                    dir_name:str=None,
                    ):
    if dir_name not in mapper.keys():
        raise ValueError(f'Choose processor from the the following: {mapper.keys()}')

    if dir_name == 'extract_dir':
        processor = NBProcessor(nb_path, partial(extract_dir, dir=dir))
        processor.process()
        return
    elif dir_name == 'no_dir_and_dir':
        processor = NBProcessor(nb_path, partial(no_dir_and_dir, dir=dir))
        processor.process()
        return

    processor = NBProcessor(nb_path, mapper[dir_name])
    processor.process()