def format_size_argument()

in scripts/gen_wrappers.py [0:0]

• 8 lines of code
• 4 McCabe index (conditional complexity)

def format_size_argument(model_var_name, shape_def):
  if isinstance(shape_def, str):
    if shape_def.startswith('n'):
      return f'{model_var_name}.{shape_def}'
    m = re.match(r'(\d+)\s*\*\s*(n[a-zA-Z]*)', shape_def)
    if m:
      return f'{m.group(1)}*{model_var_name}.{m.group(2)}'

  return shape_def