totrans

I have a suggested rewrite for this function. The main changes are:

1. Variable names no longer have to be specified. The letters "x" and "y" are always interpreted as the abscissa and ordinate, and "z" is always equal to the complex number x+iy.

2. It is now possible for the given function expr to take one complex argument and return two real ones, or to take two real argument and return one complex one.

def mytotrans(expr, globals=None, locals=None):
  """Converts to a coordinate transformation (a function that accepts
two arguments and returns two values).

expr       required                  a string expression or a function
                                     of two real or one complex value
globals    default=None              dict of global variables
locals     default=None              dict of local variables

If expr is a function, it must accept either two real or one complex
argument, and must return either one or two values.

If expr is a string, it must be a function involving the variables "x"
and "y".  It may also involve "z", which is defined to be complex(x,y),
and may use any functions in the math or cmath modules.

Functions in cmath take priority -- use math.funcname to insist on the
real version.  For example, "sqrt(x)" gives 0 + 1j when x is -1, but
"math.sqrt(x)" returns an exception.
"""

  def split(z):
    try:
      return (z.real, z.imag)
    except:
      return z
  
  if callable(expr):
    if expr.func_code.co_argcount == 1:
        thefunc = lambda x, y: expr( complex(x,y) )
    elif  expr.func_code.co_argcount == 2:
        thefunc = lambda x, y: expr(x, y)
    else:
        raise TypeError, "must be a function of 2 or 1 variables"
    output = lambda x, y: split( thefunc(x,y) )
    return output

  else:
    g = dict(math.__dict__)
    g.update(cmath.__dict__)
    g.update( [ ("math." + s, f) for s,f in math.__dict__.iteritems() ] )
    g.update( [ ("cmath." + s, f) for s,f in cmath.__dict__.iteritems() ] )
    
    if globals != None: g.update(globals)
    thefunc = eval("lambda x, y, z: (%s)" % (expr), g, locals)
    output = lambda x, y: split( thefunc(x, y, complex(x,y)) )
    output.func_name = "totrans('%s')" % (expr)
    return output

I'm no longer convinced that this code works. I'll let you know more later.

I liked the idea of dynamically determining if the string expression is in terms of x,y or in terms of z. Here's a way that you can construct an "if" statement to switch between the two cases:

varnames = compile(expr, "", "eval").co_names
if "x" in varnames and "y" in varnames:
    it's real
elif "z" in varnames:
    it's complex

Then you can use this to choose math or cmath. I'd prefer the cleanliness of one or the other, but not both. I'm sure that the math libraries are faster to compute than the cmath libraries, and the user will probably use real numbers 95% of the time.

All right, here is another attempt:

def trans(expr, vars = None, globals = {}, locals = {}):
  """Converts to a coordinate transformation (a function that accepts
two arguments and returns two values).

expr       required                  a string expression or a function
                                     of two real or one complex value
vars       default=None              independent variable names;
                                     if vars is None then the function
                                     attempts to guess
globals    default={}                dict of global variables
locals     default={}                dict of local variables
"""
  if callable(expr):
    if expr.func_code.co_argcount == 2:
      return expr

    elif expr.func_code.co_argcount == 1:
      split = lambda z: (z.real, z.imag)
      output = lambda x, y: split(expr(x + y*1j))
      output.func_name = expr.func_name
      return output

    else:
      raise TypeError, "must be a function of 1 or 2 variables"

  if vars is None:
    varnames = compile(expr, "", "eval").co_names
    if 'x' in varnames and 'y' in varnames:
      vars = ('x', 'y')
    elif 'z' in varnames:
      vars = ['z']
    else:   # function will attempt to guess variable names
      varnames = [x for x in varnames if (x not in math.__dict__
                                      and x not in cmath.__dict__
                                      and x not in __builtins__.__dict__
                                      and x not in globals
                                      and x not in locals)]    # remove known names
      if len(varnames) == 1:
        vars = varnames
      elif len(varnames) == 2:
        vars = varnames
        vars.sort()  # alphabetical order assumed
      else:
        raise SyntaxError, 'cannot guess variable names'

  elif isinstance(vars, str):
    if ',' in vars:
      vars = vars.split(',')
      if len(vars) == 2:
        vars[0].strip(); vars[1].strip()
      else:
        raise TypeError, 'vars must have 1 or 2 elements'
    else:
      vars = [vars.strip()]

  if len(vars) == 2:    # real case
    g = math.__dict__
    if globals != None: g.update(globals)
    output = eval("lambda %s, %s: (%s)" % (vars[0], vars[1], expr), g, locals)
    output.func_name = "(%s, %s) -> (%s)" % (vars[0], vars[1], expr)
    return output

  elif len(vars) == 1:  # complex case
    g = cmath.__dict__
    if globals != None: g.update(globals)
    output = eval("lambda %s: (%s)" % (vars[0], expr), g, locals)
    split = lambda z: (z.real, z.imag)
    output2 = lambda x, y: split(output(x + y*1j))
    output2.func_name = "%s -> %s" % (vars[0], expr)
    return output2

  else:
    raise TypeError, "vars must have 1 or 2 elements"

Notable features: 1. I think "trans" is a better name than "totrans". It's shorter, and it agrees with notations like int('3700'). 2. I used the code you suggested to make it interpret expressions involving 'x' and 'y' or 'z' correctly. 3. Otherwise it attempts to guess the variables names. I tested the algorithm and and it works reasonably well:

>>> trans('r * cos(theta), r * sin(theta)')
<function (r, theta) -> (r * cos(theta), r*sin(theta)) at 0x020BABB0>
>>> trans('u**2 + v**2, 2*u*v')
<function (u, v) -> (u**2 + v**2, 2*u*v) at 0x020BA4F0>
>>> trans('x2-x1,x2+x1')
<function (x1, x2) -> (x2-x1,x2+x1) at 0x020BA1B0>
>>> trans('(w+1)/(w-1)')
<function w -> (w+1)/(w-1) at 0x020DD770>

4. It is also tolerant towards strings with commas entered as variable lists:

>>> trans('phi + chi, phi * chi', 'phi, chi')
<function (phi,  chi) -> (phi + chi, phi * chi) at 0x0138E0F0>

What do you think?

I like it!

Go ahead and upload it.