yprov4dask/parser.py at NextRelease · Yousaf95/yprov4dask · 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
115
116
117
import tree_sitter_python as tspython
from tree_sitter import Language, Node, Parser, Query, QueryCursor

language = Language(tspython.language())
parser = Parser(language)

triggers = """
  "compute"
  "submit"
  "get"
  "persist"
  "fire_and_forget"
  "store"
  "to_zarr"
  "to_hdf5"
  "to_tiledb"
  "plot"
  "open_dataset"
"""
triggers_set = set(["compute", "submit", "get", "persist", "fire_and_forget",
  "store", "to_zarr", "to_hdf5", "to_tiledb", "plot", "hvplot"])
conditionals = set(['store', 'to_zarr', 'to_hdf5', 'to_tiledb'])
valid_hints = """
  "# Jupyter:ignore"
  "# Jupyter:track"
"""

code = """
def save_ds(delayed):
    # Jupyter:ignore
    delayed.parsist()

flag = False
# Jupyter:ignore
delayed = ds.to_zarr('output.zarr', compute = flag)
save_ds(delayed) # Jupyter:track
ds = xr.tutorial[444].ciao().ciao.open_dataset()
air = ds['air'].persist()
fire_and_forget()
season = pn.widgets.Select(name='Season', options=list(seasonal_climatology.season.values))
# Jupyter:ignore
seasonal_climatology.interactive().sel(season=season).hvplot(cmap='seismic')
# Jupyter:track
f.air[2][1](1, 2)
"""
tree = parser.parse(bytes(code, 'utf-8'))

call_query = Query(language,
  f"""[
    (call
      [
        (
          (identifier) @id
          (argument_list) @args
        )
        (
          (attribute
            [(call) (identifier) (subscript)]*
            "."
            (identifier) @id
          ) (argument_list) @args
        )
        (
          (subscript) @id
          (argument_list) @args
        )
      ]
    )
    ((comment) @hint
      (#any-of? @hint {valid_hints}))
  ]"""
)
call_cursor = QueryCursor(call_query)
args_query = Query(language,
  """(keyword_argument
    name: ((identifier) @name
      (#eq? @name "compute"))
    "="
    value: (false) @value
  )"""
)
args_cursor = QueryCursor(args_query)
hint_query = Query(language,
  f"""((comment) @hint
    (#any-of? @hint {valid_hints}))
  """
)
hint_cursor = QueryCursor(hint_query)

#comments = hint_cursor.captures(tree.root_node)
calls = call_cursor.captures(tree.root_node)
comments = calls.get('hint', [])
print('comments')
for comm in comments:
  print(comm.text.decode('utf-8'))

ids = calls['id']
args = calls['args']
hints = calls['hint']

ids = sorted(ids, key=lambda node: node.start_byte)
args = sorted(args, key=lambda node: node.start_byte)
hints = sorted(hints, key=lambda node: node.start_byte)
nodes = zip(ids, args)
for id, arg in nodes:
  keep = id.text.decode('utf-8') in triggers_set
  if len(hints) > 0:
    if id.start_byte > hints[0].start_byte:
      hint_node = hints[0]
      text = hint_node.text.decode('utf-8')
      if keep and text == '# Jupyter:ignore':
        keep = False
        hints.remove(hint_node)
      elif not keep and text == '# Jupyter:track':
        keep = True
        hints.remove(hint_node)
  print(f"{id.text} {arg.text} {'Track' if keep else 'Ignore'}")