Enable valid Python to be a config.gin file, so code analysis and syntax highlighting works

python_syntax_for_gin.py

def preprocess_config(s: str):
  """Remove imports from a string representation of a python file"""
  # We assume that imports are not multi-line.
  lines = s.splitlines()
  out_lines = []

  for line in lines:
    # Skip lines with import in them:
    if 'import' in line:
      continue
    # We add "@" to each symbol after the = sign.
    # For example:
    # `MlpTowerFactory.nonlinearity = softplus``
    # becomes:
    # `MlpTowerFactory.nonlinearity = @softplus`
    # However, this does not happen if a % symbol is there.
    # It also doesn't happen if the symbol (softplus in this example)
    # starts with a number.
    if "=" in line and "%" not in line and not line.startswith("#"):
      tokens = line.replace(" ", "").split("=")
      assert len(tokens) == 2
      if tokens[-1][0].isalpha():
        line = f"{tokens[0]} = @{tokens[-1]}"

    out_lines.append(line)

  return '\n'.join(out_lines)

e.g.,

config.gin

from mypackage.file1 import object
from mypackage.file2 import function

object.create.arg = function
object.create.arg2 = 2.0  # These still work!

setting = 1.0
object.create.arg3 = %SETTING  # These are left untouched

Now when you hover over function, Python analysis tools can correctly find the source of it, rather than the @ destroying the reference.

Simply load a config file, run preprocess_config on it, then pass it to gin.