1  from typing import Dict, Any, Callable, Iterable
 2  from eth2spec.debug.encode import encode
 3  
 4  
 5  def spectest(description: str = None):
 6      def runner(fn):
 7          # this wraps the function, to hide that the function actually is yielding data, instead of returning once.
 8          def entry(*args, **kw):
 9              # check generator mode, may be None/else.
10              # "pop" removes it, so it is not passed to the inner function.
11              if kw.pop('generator_mode', False) is True:
12                  out = {}
13                  if description is None:
14                      # fall back on function name for test description
15                      name = fn.__name__
16                      if name.startswith('test_'):
17                          name = name[5:]
18                      out['description'] = name
19                  else:
20                      # description can be explicit
21                      out['description'] = description
22                  has_contents = False
23                  # put all generated data into a dict.
24                  for data in fn(*args, **kw):
25                      has_contents = True
26                      # If there is a type argument, encode it as that type.
27                      if len(data) == 3:
28                          (key, value, typ) = data
29                          out[key] = encode(value, typ)
30                      else:
31                          # Otherwise, try to infer the type, but keep it as-is if it's not a SSZ container.
32                          (key, value) = data
33                          if hasattr(value.__class__, 'fields'):
34                              out[key] = encode(value, value.__class__)
35                          else:
36                              out[key] = value
37                  if has_contents:
38                      return out
39                  else:
40                      return None
41              else:
42                  # just complete the function, ignore all yielded data, we are not using it
43                  for _ in fn(*args, **kw):
44                      continue
45                  return None
46          return entry
47      return runner
48  
49  
50  def with_tags(tags: Dict[str, Any]):
51      """
52      Decorator factory, adds tags (key, value) pairs to the output of the function.
53      Useful to build test-vector annotations with.
54      This decorator is applied after the ``spectest`` decorator is applied.
55      :param tags: dict of tags
56      :return: Decorator.
57      """
58      def runner(fn):
59          def entry(*args, **kw):
60              fn_out = fn(*args, **kw)
61              # do not add tags if the function is not returning a dict at all (i.e. not in generator mode)
62              if fn_out is None:
63                  return None
64              return {**tags, **fn_out}
65          return entry
66      return runner
67  
68  
69  def with_args(create_args: Callable[[], Iterable[Any]]):
70      """
71      Decorator factory, adds given extra arguments to the decorated function.
72      :param create_args: function to create arguments with.
73      :return: Decorator.
74      """
75      def runner(fn):
76          # this wraps the function, to hide that the function actually yielding data.
77          def entry(*args, **kw):
78              return fn(*(list(create_args()) + list(args)), **kw)
79          return entry
80      return runner