1*9c5db199SXin Li# Lint as: python2, python3 2*9c5db199SXin Li# pylint: disable-msg=C0111 3*9c5db199SXin Li# Copyright 2008 Google Inc. Released under the GPL v2 4*9c5db199SXin Li 5*9c5db199SXin Lifrom __future__ import absolute_import 6*9c5db199SXin Lifrom __future__ import division 7*9c5db199SXin Lifrom __future__ import print_function 8*9c5db199SXin Li 9*9c5db199SXin Liimport ast 10*9c5db199SXin Liimport logging 11*9c5db199SXin Liimport textwrap 12*9c5db199SXin Liimport re 13*9c5db199SXin Liimport six 14*9c5db199SXin Li 15*9c5db199SXin Lifrom autotest_lib.client.common_lib import autotest_enum 16*9c5db199SXin Lifrom autotest_lib.client.common_lib import global_config 17*9c5db199SXin Lifrom autotest_lib.client.common_lib import priorities 18*9c5db199SXin Li 19*9c5db199SXin Li 20*9c5db199SXin LiREQUIRED_VARS = set(['author', 'doc', 'name', 'time', 'test_type']) 21*9c5db199SXin LiOBSOLETE_VARS = set(['experimental']) 22*9c5db199SXin Li 23*9c5db199SXin LiCONTROL_TYPE = autotest_enum.AutotestEnum('Server', 'Client', start_value=1) 24*9c5db199SXin LiCONTROL_TYPE_NAMES = autotest_enum.AutotestEnum(*CONTROL_TYPE.names, 25*9c5db199SXin Li string_values=True) 26*9c5db199SXin Li 27*9c5db199SXin Li_SUITE_ATTRIBUTE_PREFIX = 'suite:' 28*9c5db199SXin Li 29*9c5db199SXin LiCONFIG = global_config.global_config 30*9c5db199SXin Li 31*9c5db199SXin Li# Default maximum test result size in kB. 32*9c5db199SXin LiDEFAULT_MAX_RESULT_SIZE_KB = CONFIG.get_config_value( 33*9c5db199SXin Li 'AUTOSERV', 'default_max_result_size_KB', type=int, default=20000) 34*9c5db199SXin Li 35*9c5db199SXin Li 36*9c5db199SXin Liclass ControlVariableException(Exception): 37*9c5db199SXin Li pass 38*9c5db199SXin Li 39*9c5db199SXin Lidef _validate_control_file_fields(control_file_path, control_file_vars, 40*9c5db199SXin Li raise_warnings): 41*9c5db199SXin Li """Validate the given set of variables from a control file. 42*9c5db199SXin Li 43*9c5db199SXin Li @param control_file_path: string path of the control file these were 44*9c5db199SXin Li loaded from. 45*9c5db199SXin Li @param control_file_vars: dict of variables set in a control file. 46*9c5db199SXin Li @param raise_warnings: True iff we should raise on invalid variables. 47*9c5db199SXin Li 48*9c5db199SXin Li """ 49*9c5db199SXin Li diff = REQUIRED_VARS - set(control_file_vars) 50*9c5db199SXin Li if diff: 51*9c5db199SXin Li warning = ('WARNING: Not all required control ' 52*9c5db199SXin Li 'variables were specified in %s. Please define ' 53*9c5db199SXin Li '%s.') % (control_file_path, ', '.join(diff)) 54*9c5db199SXin Li if raise_warnings: 55*9c5db199SXin Li raise ControlVariableException(warning) 56*9c5db199SXin Li print(textwrap.wrap(warning, 80)) 57*9c5db199SXin Li 58*9c5db199SXin Li obsolete = OBSOLETE_VARS & set(control_file_vars) 59*9c5db199SXin Li if obsolete: 60*9c5db199SXin Li warning = ('WARNING: Obsolete variables were ' 61*9c5db199SXin Li 'specified in %s. Please remove ' 62*9c5db199SXin Li '%s.') % (control_file_path, ', '.join(obsolete)) 63*9c5db199SXin Li if raise_warnings: 64*9c5db199SXin Li raise ControlVariableException(warning) 65*9c5db199SXin Li print(textwrap.wrap(warning, 80)) 66*9c5db199SXin Li 67*9c5db199SXin Li 68*9c5db199SXin Liclass ControlData(object): 69*9c5db199SXin Li # Available TIME settings in control file, the list must be in lower case 70*9c5db199SXin Li # and in ascending order, test running faster comes first. 71*9c5db199SXin Li TEST_TIME_LIST = ['fast', 'short', 'medium', 'long', 'lengthy'] 72*9c5db199SXin Li TEST_TIME = autotest_enum.AutotestEnum(*TEST_TIME_LIST, 73*9c5db199SXin Li string_values=False) 74*9c5db199SXin Li 75*9c5db199SXin Li @staticmethod 76*9c5db199SXin Li def get_test_time_index(time): 77*9c5db199SXin Li """ 78*9c5db199SXin Li Get the order of estimated test time, based on the TIME setting in 79*9c5db199SXin Li Control file. Faster test gets a lower index number. 80*9c5db199SXin Li """ 81*9c5db199SXin Li try: 82*9c5db199SXin Li return ControlData.TEST_TIME.get_value(time.lower()) 83*9c5db199SXin Li except AttributeError: 84*9c5db199SXin Li # Raise exception if time value is not a valid TIME setting. 85*9c5db199SXin Li error_msg = '%s is not a valid TIME.' % time 86*9c5db199SXin Li logging.error(error_msg) 87*9c5db199SXin Li raise ControlVariableException(error_msg) 88*9c5db199SXin Li 89*9c5db199SXin Li 90*9c5db199SXin Li def __init__(self, vars, path, raise_warnings=False): 91*9c5db199SXin Li # Defaults 92*9c5db199SXin Li self.path = path 93*9c5db199SXin Li self.dependencies = set() 94*9c5db199SXin Li # TODO(jrbarnette): This should be removed once outside 95*9c5db199SXin Li # code that uses can be changed. 96*9c5db199SXin Li self.experimental = False 97*9c5db199SXin Li self.run_verify = True 98*9c5db199SXin Li self.sync_count = 1 99*9c5db199SXin Li self.test_parameters = set() 100*9c5db199SXin Li self.test_category = '' 101*9c5db199SXin Li self.test_class = '' 102*9c5db199SXin Li self.job_retries = 0 103*9c5db199SXin Li # Default to require server-side package. Unless require_ssp is 104*9c5db199SXin Li # explicitly set to False, server-side package will be used for the 105*9c5db199SXin Li # job. 106*9c5db199SXin Li self.require_ssp = None 107*9c5db199SXin Li self.attributes = set() 108*9c5db199SXin Li self.max_result_size_KB = DEFAULT_MAX_RESULT_SIZE_KB 109*9c5db199SXin Li self.priority = priorities.Priority.DEFAULT 110*9c5db199SXin Li self.extended_timeout = None 111*9c5db199SXin Li self.fast = True 112*9c5db199SXin Li # This will only be honored via `test_that`, and not in lab (for now). 113*9c5db199SXin Li self.py_version = None 114*9c5db199SXin Li 115*9c5db199SXin Li _validate_control_file_fields(self.path, vars, raise_warnings) 116*9c5db199SXin Li 117*9c5db199SXin Li for key, val in six.iteritems(vars): 118*9c5db199SXin Li try: 119*9c5db199SXin Li self.set_attr(key, val, raise_warnings) 120*9c5db199SXin Li except Exception as e: 121*9c5db199SXin Li if raise_warnings: 122*9c5db199SXin Li raise 123*9c5db199SXin Li print('WARNING: %s; skipping' % e) 124*9c5db199SXin Li 125*9c5db199SXin Li self._patch_up_suites_from_attributes() 126*9c5db199SXin Li 127*9c5db199SXin Li 128*9c5db199SXin Li @property 129*9c5db199SXin Li def suite_tag_parts(self): 130*9c5db199SXin Li """Return the part strings of the test's suite tag.""" 131*9c5db199SXin Li if hasattr(self, 'suite'): 132*9c5db199SXin Li return [part.strip() for part in self.suite.split(',')] 133*9c5db199SXin Li else: 134*9c5db199SXin Li return [] 135*9c5db199SXin Li 136*9c5db199SXin Li 137*9c5db199SXin Li def set_attr(self, attr, val, raise_warnings=False): 138*9c5db199SXin Li attr = attr.lower() 139*9c5db199SXin Li try: 140*9c5db199SXin Li set_fn = getattr(self, 'set_%s' % attr) 141*9c5db199SXin Li set_fn(val) 142*9c5db199SXin Li except AttributeError: 143*9c5db199SXin Li # This must not be a variable we care about 144*9c5db199SXin Li pass 145*9c5db199SXin Li 146*9c5db199SXin Li 147*9c5db199SXin Li def _patch_up_suites_from_attributes(self): 148*9c5db199SXin Li """Patch up the set of suites this test is part of. 149*9c5db199SXin Li 150*9c5db199SXin Li Legacy builds will not have an appropriate ATTRIBUTES field set. 151*9c5db199SXin Li Take the union of suites specified via ATTRIBUTES and suites specified 152*9c5db199SXin Li via SUITE. 153*9c5db199SXin Li 154*9c5db199SXin Li SUITE used to be its own variable, but now suites are taken only from 155*9c5db199SXin Li the attributes. 156*9c5db199SXin Li 157*9c5db199SXin Li """ 158*9c5db199SXin Li 159*9c5db199SXin Li suite_names = set() 160*9c5db199SXin Li # Extract any suites we know ourselves to be in based on the SUITE 161*9c5db199SXin Li # line. This line is deprecated, but control files in old builds will 162*9c5db199SXin Li # still have it. 163*9c5db199SXin Li if hasattr(self, 'suite'): 164*9c5db199SXin Li existing_suites = self.suite.split(',') 165*9c5db199SXin Li existing_suites = [name.strip() for name in existing_suites] 166*9c5db199SXin Li existing_suites = [name for name in existing_suites if name] 167*9c5db199SXin Li suite_names.update(existing_suites) 168*9c5db199SXin Li 169*9c5db199SXin Li # Figure out if our attributes mention any suites. 170*9c5db199SXin Li for attribute in self.attributes: 171*9c5db199SXin Li if not attribute.startswith(_SUITE_ATTRIBUTE_PREFIX): 172*9c5db199SXin Li continue 173*9c5db199SXin Li suite_name = attribute[len(_SUITE_ATTRIBUTE_PREFIX):] 174*9c5db199SXin Li suite_names.add(suite_name) 175*9c5db199SXin Li 176*9c5db199SXin Li # Rebuild the suite field if necessary. 177*9c5db199SXin Li if suite_names: 178*9c5db199SXin Li self.set_suite(','.join(sorted(list(suite_names)))) 179*9c5db199SXin Li 180*9c5db199SXin Li 181*9c5db199SXin Li def _set_string(self, attr, val): 182*9c5db199SXin Li val = str(val) 183*9c5db199SXin Li setattr(self, attr, val) 184*9c5db199SXin Li 185*9c5db199SXin Li 186*9c5db199SXin Li def _set_option(self, attr, val, options): 187*9c5db199SXin Li val = str(val) 188*9c5db199SXin Li if val.lower() not in [x.lower() for x in options]: 189*9c5db199SXin Li raise ValueError("%s must be one of the following " 190*9c5db199SXin Li "options: %s" % (attr, 191*9c5db199SXin Li ', '.join(options))) 192*9c5db199SXin Li setattr(self, attr, val) 193*9c5db199SXin Li 194*9c5db199SXin Li 195*9c5db199SXin Li def _set_bool(self, attr, val): 196*9c5db199SXin Li val = str(val).lower() 197*9c5db199SXin Li if val == "false": 198*9c5db199SXin Li val = False 199*9c5db199SXin Li elif val == "true": 200*9c5db199SXin Li val = True 201*9c5db199SXin Li else: 202*9c5db199SXin Li msg = "%s must be either true or false" % attr 203*9c5db199SXin Li raise ValueError(msg) 204*9c5db199SXin Li setattr(self, attr, val) 205*9c5db199SXin Li 206*9c5db199SXin Li 207*9c5db199SXin Li def _set_int(self, attr, val, min=None, max=None): 208*9c5db199SXin Li val = int(val) 209*9c5db199SXin Li if min is not None and min > val: 210*9c5db199SXin Li raise ValueError("%s is %d, which is below the " 211*9c5db199SXin Li "minimum of %d" % (attr, val, min)) 212*9c5db199SXin Li if max is not None and max < val: 213*9c5db199SXin Li raise ValueError("%s is %d, which is above the " 214*9c5db199SXin Li "maximum of %d" % (attr, val, max)) 215*9c5db199SXin Li setattr(self, attr, val) 216*9c5db199SXin Li 217*9c5db199SXin Li 218*9c5db199SXin Li def _set_set(self, attr, val): 219*9c5db199SXin Li val = str(val) 220*9c5db199SXin Li items = [x.strip() for x in val.split(',') if x.strip()] 221*9c5db199SXin Li setattr(self, attr, set(items)) 222*9c5db199SXin Li 223*9c5db199SXin Li 224*9c5db199SXin Li def set_author(self, val): 225*9c5db199SXin Li self._set_string('author', val) 226*9c5db199SXin Li 227*9c5db199SXin Li 228*9c5db199SXin Li def set_dependencies(self, val): 229*9c5db199SXin Li self._set_set('dependencies', val) 230*9c5db199SXin Li 231*9c5db199SXin Li 232*9c5db199SXin Li def set_doc(self, val): 233*9c5db199SXin Li self._set_string('doc', val) 234*9c5db199SXin Li 235*9c5db199SXin Li 236*9c5db199SXin Li def set_name(self, val): 237*9c5db199SXin Li self._set_string('name', val) 238*9c5db199SXin Li 239*9c5db199SXin Li 240*9c5db199SXin Li def set_run_verify(self, val): 241*9c5db199SXin Li self._set_bool('run_verify', val) 242*9c5db199SXin Li 243*9c5db199SXin Li 244*9c5db199SXin Li def set_sync_count(self, val): 245*9c5db199SXin Li self._set_int('sync_count', val, min=1) 246*9c5db199SXin Li 247*9c5db199SXin Li 248*9c5db199SXin Li def set_suite(self, val): 249*9c5db199SXin Li self._set_string('suite', val) 250*9c5db199SXin Li 251*9c5db199SXin Li 252*9c5db199SXin Li def set_time(self, val): 253*9c5db199SXin Li self._set_option('time', val, ControlData.TEST_TIME_LIST) 254*9c5db199SXin Li 255*9c5db199SXin Li 256*9c5db199SXin Li def set_test_class(self, val): 257*9c5db199SXin Li self._set_string('test_class', val.lower()) 258*9c5db199SXin Li 259*9c5db199SXin Li 260*9c5db199SXin Li def set_test_category(self, val): 261*9c5db199SXin Li self._set_string('test_category', val.lower()) 262*9c5db199SXin Li 263*9c5db199SXin Li 264*9c5db199SXin Li def set_test_type(self, val): 265*9c5db199SXin Li self._set_option('test_type', val, list(CONTROL_TYPE.names)) 266*9c5db199SXin Li 267*9c5db199SXin Li 268*9c5db199SXin Li def set_test_parameters(self, val): 269*9c5db199SXin Li self._set_set('test_parameters', val) 270*9c5db199SXin Li 271*9c5db199SXin Li 272*9c5db199SXin Li def set_job_retries(self, val): 273*9c5db199SXin Li self._set_int('job_retries', val) 274*9c5db199SXin Li 275*9c5db199SXin Li 276*9c5db199SXin Li def set_bug_template(self, val): 277*9c5db199SXin Li if type(val) == dict: 278*9c5db199SXin Li setattr(self, 'bug_template', val) 279*9c5db199SXin Li 280*9c5db199SXin Li 281*9c5db199SXin Li def set_require_ssp(self, val): 282*9c5db199SXin Li self._set_bool('require_ssp', val) 283*9c5db199SXin Li 284*9c5db199SXin Li 285*9c5db199SXin Li def set_build(self, val): 286*9c5db199SXin Li self._set_string('build', val) 287*9c5db199SXin Li 288*9c5db199SXin Li 289*9c5db199SXin Li def set_builds(self, val): 290*9c5db199SXin Li if type(val) == dict: 291*9c5db199SXin Li setattr(self, 'builds', val) 292*9c5db199SXin Li 293*9c5db199SXin Li def set_max_result_size_kb(self, val): 294*9c5db199SXin Li self._set_int('max_result_size_KB', val) 295*9c5db199SXin Li 296*9c5db199SXin Li def set_priority(self, val): 297*9c5db199SXin Li self._set_int('priority', val) 298*9c5db199SXin Li 299*9c5db199SXin Li def set_fast(self, val): 300*9c5db199SXin Li self._set_bool('fast', val) 301*9c5db199SXin Li 302*9c5db199SXin Li def set_update_type(self, val): 303*9c5db199SXin Li self._set_string('update_type', val) 304*9c5db199SXin Li 305*9c5db199SXin Li def set_source_release(self, val): 306*9c5db199SXin Li self._set_string('source_release', val) 307*9c5db199SXin Li 308*9c5db199SXin Li def set_target_release(self, val): 309*9c5db199SXin Li self._set_string('target_release', val) 310*9c5db199SXin Li 311*9c5db199SXin Li def set_target_payload_uri(self, val): 312*9c5db199SXin Li self._set_string('target_payload_uri', val) 313*9c5db199SXin Li 314*9c5db199SXin Li def set_source_payload_uri(self, val): 315*9c5db199SXin Li self._set_string('source_payload_uri', val) 316*9c5db199SXin Li 317*9c5db199SXin Li def set_source_archive_uri(self, val): 318*9c5db199SXin Li self._set_string('source_archive_uri', val) 319*9c5db199SXin Li 320*9c5db199SXin Li def set_attributes(self, val): 321*9c5db199SXin Li self._set_set('attributes', val) 322*9c5db199SXin Li 323*9c5db199SXin Li def set_extended_timeout(self, val): 324*9c5db199SXin Li """In seconds.""" 325*9c5db199SXin Li self._set_int('extended_timeout', val) 326*9c5db199SXin Li 327*9c5db199SXin Li def set_py_version(self, val): 328*9c5db199SXin Li """In majors, ie: 2 or 3.""" 329*9c5db199SXin Li self._set_int('py_version', val) 330*9c5db199SXin Li 331*9c5db199SXin Li 332*9c5db199SXin Lidef _extract_const(expr): 333*9c5db199SXin Li assert (expr.__class__ == ast.Str) 334*9c5db199SXin Li if six.PY2: 335*9c5db199SXin Li assert (expr.s.__class__ in (str, int, float, unicode)) 336*9c5db199SXin Li else: 337*9c5db199SXin Li assert (expr.s.__class__ in (str, int, float)) 338*9c5db199SXin Li return str(expr.s).strip() 339*9c5db199SXin Li 340*9c5db199SXin Li 341*9c5db199SXin Lidef _extract_dict(expr): 342*9c5db199SXin Li assert (expr.__class__ == ast.Dict) 343*9c5db199SXin Li assert (expr.keys.__class__ == list) 344*9c5db199SXin Li cf_dict = {} 345*9c5db199SXin Li for key, value in zip(expr.keys, expr.values): 346*9c5db199SXin Li try: 347*9c5db199SXin Li key = _extract_const(key) 348*9c5db199SXin Li val = _extract_expression(value) 349*9c5db199SXin Li except (AssertionError, ValueError): 350*9c5db199SXin Li pass 351*9c5db199SXin Li else: 352*9c5db199SXin Li cf_dict[key] = val 353*9c5db199SXin Li return cf_dict 354*9c5db199SXin Li 355*9c5db199SXin Li 356*9c5db199SXin Lidef _extract_list(expr): 357*9c5db199SXin Li assert (expr.__class__ == ast.List) 358*9c5db199SXin Li list_values = [] 359*9c5db199SXin Li for value in expr.elts: 360*9c5db199SXin Li try: 361*9c5db199SXin Li list_values.append(_extract_expression(value)) 362*9c5db199SXin Li except (AssertionError, ValueError): 363*9c5db199SXin Li pass 364*9c5db199SXin Li return list_values 365*9c5db199SXin Li 366*9c5db199SXin Li 367*9c5db199SXin Lidef _extract_name(expr): 368*9c5db199SXin Li assert (expr.__class__ == ast.Name) 369*9c5db199SXin Li assert (expr.id in ('False', 'True', 'None')) 370*9c5db199SXin Li return str(expr.id) 371*9c5db199SXin Li 372*9c5db199SXin Li 373*9c5db199SXin Lidef _extract_expression(expr): 374*9c5db199SXin Li if expr.__class__ == ast.Str: 375*9c5db199SXin Li return _extract_const(expr) 376*9c5db199SXin Li if expr.__class__ == ast.Name: 377*9c5db199SXin Li return _extract_name(expr) 378*9c5db199SXin Li if expr.__class__ == ast.Dict: 379*9c5db199SXin Li return _extract_dict(expr) 380*9c5db199SXin Li if expr.__class__ == ast.List: 381*9c5db199SXin Li return _extract_list(expr) 382*9c5db199SXin Li if expr.__class__ == ast.Num: 383*9c5db199SXin Li return expr.n 384*9c5db199SXin Li if six.PY3 and expr.__class__ == ast.NameConstant: 385*9c5db199SXin Li return expr.value 386*9c5db199SXin Li if six.PY3 and expr.__class__ == ast.Constant: 387*9c5db199SXin Li try: 388*9c5db199SXin Li return expr.value.strip() 389*9c5db199SXin Li except Exception: 390*9c5db199SXin Li return expr.value 391*9c5db199SXin Li raise ValueError('Unknown rval %s' % expr) 392*9c5db199SXin Li 393*9c5db199SXin Li 394*9c5db199SXin Lidef _extract_assignment(n): 395*9c5db199SXin Li assert (n.__class__ == ast.Assign) 396*9c5db199SXin Li assert (len(n.targets) == 1) 397*9c5db199SXin Li assert (n.targets[0].__class__ == ast.Name) 398*9c5db199SXin Li val = _extract_expression(n.value) 399*9c5db199SXin Li key = n.targets[0].id.lower() 400*9c5db199SXin Li return (key, val) 401*9c5db199SXin Li 402*9c5db199SXin Li 403*9c5db199SXin Lidef parse_control_string(control, raise_warnings=False, path=''): 404*9c5db199SXin Li """Parse a control file from a string. 405*9c5db199SXin Li 406*9c5db199SXin Li @param control: string containing the text of a control file. 407*9c5db199SXin Li @param raise_warnings: True iff ControlData should raise an error on 408*9c5db199SXin Li warnings about control file contents. 409*9c5db199SXin Li @param path: string path to the control file. 410*9c5db199SXin Li 411*9c5db199SXin Li """ 412*9c5db199SXin Li try: 413*9c5db199SXin Li mod = ast.parse(control) 414*9c5db199SXin Li except SyntaxError as e: 415*9c5db199SXin Li logging.error('Syntax error (%s) while parsing control string:', e) 416*9c5db199SXin Li lines = control.split('\n') 417*9c5db199SXin Li for n, l in enumerate(lines): 418*9c5db199SXin Li logging.error('Line %d: %s', n + 1, l) 419*9c5db199SXin Li raise ControlVariableException("Error parsing data because %s" % e) 420*9c5db199SXin Li return finish_parse(mod, path, raise_warnings) 421*9c5db199SXin Li 422*9c5db199SXin Li 423*9c5db199SXin Lidef parse_control(path, raise_warnings=False): 424*9c5db199SXin Li try: 425*9c5db199SXin Li with open(path, 'r') as r: 426*9c5db199SXin Li mod = ast.parse(r.read()) 427*9c5db199SXin Li except SyntaxError as e: 428*9c5db199SXin Li raise ControlVariableException("Error parsing %s because %s" % 429*9c5db199SXin Li (path, e)) 430*9c5db199SXin Li return finish_parse(mod, path, raise_warnings) 431*9c5db199SXin Li 432*9c5db199SXin Li 433*9c5db199SXin Lidef _try_extract_assignment(node, variables): 434*9c5db199SXin Li """Try to extract assignment from the given node. 435*9c5db199SXin Li 436*9c5db199SXin Li @param node: An Assign object. 437*9c5db199SXin Li @param variables: Dictionary to store the parsed assignments. 438*9c5db199SXin Li """ 439*9c5db199SXin Li try: 440*9c5db199SXin Li key, val = _extract_assignment(node) 441*9c5db199SXin Li variables[key] = val 442*9c5db199SXin Li except (AssertionError, ValueError) as e: 443*9c5db199SXin Li pass 444*9c5db199SXin Li 445*9c5db199SXin Li 446*9c5db199SXin Lidef finish_parse(mod, path, raise_warnings): 447*9c5db199SXin Li assert (mod.__class__ == ast.Module) 448*9c5db199SXin Li assert (mod.body.__class__ == list) 449*9c5db199SXin Li 450*9c5db199SXin Li variables = {} 451*9c5db199SXin Li injection_variables = {} 452*9c5db199SXin Li for n in mod.body: 453*9c5db199SXin Li if (n.__class__ == ast.FunctionDef and re.match('step\d+', n.name)): 454*9c5db199SXin Li vars_in_step = {} 455*9c5db199SXin Li for sub_node in n.body: 456*9c5db199SXin Li _try_extract_assignment(sub_node, vars_in_step) 457*9c5db199SXin Li if vars_in_step: 458*9c5db199SXin Li # Empty the vars collection so assignments from multiple steps 459*9c5db199SXin Li # won't be mixed. 460*9c5db199SXin Li variables.clear() 461*9c5db199SXin Li variables.update(vars_in_step) 462*9c5db199SXin Li else: 463*9c5db199SXin Li _try_extract_assignment(n, injection_variables) 464*9c5db199SXin Li 465*9c5db199SXin Li variables.update(injection_variables) 466*9c5db199SXin Li return ControlData(variables, path, raise_warnings) 467