import sys
import os
import imp
import unittest
import inspect

from test.test_support import TESTFN, run_unittest

def sourcerange(source, top, bottom):
    lines = source.split("\n")
    return "\n".join(lines[top-1:bottom]) + "\n"

onelinesource = r"""# line 1
oll = lambda m: m

# line 4
tll = lambda g: g and \
g and \
g

# line 9
tlli = lambda d: d and \
    d

# line 13
def onelinefunc(): pass

# line 16
def manyargs(arg1, arg2,
arg3, arg4): pass

# line 20
def twolinefunc(m): return m and \
m

# line 24
a = [None,
     lambda x: x,
     None]

# line 29
def setfunc(func):
    globals()["anonymous"] = func
setfunc(lambda x, y: x*y)
"""

class TestOneliners(unittest.TestCase):
    
    def __init__(self, *args, **kwargs):
        unittest.TestCase.__init__(self, *args, **kwargs)
        
        self.source = onelinesource
    
    def test_oneline_lambda(self):
        # Test inspect.getsource with a one-line lambda function.
        self.assertEqual(inspect.getsource(mod.oll),
                         sourcerange(self.source, 2, 2))
    
    def test_twoline_lambda(self):
        # Test inspect.getsource with a three-line lambda function,
        # where the second and third lines are _not_ indented.
        self.assertEqual(inspect.getsource(mod.tll),
                         sourcerange(self.source, 5, 7))
    
    def test_twoline_lambda_indented(self):
        # Test inspect.getsource with a two-line lambda function, where the
        # second line _is_ indented.
        self.assertEqual(inspect.getsource(mod.tlli),
                         sourcerange(self.source, 10, 11))
    
    def test_oneline_func(self):
        # Test inspect.getsource with a regular one-line function.
        self.assertEqual(inspect.getsource(mod.onelinefunc),
                         sourcerange(self.source, 14, 14))
    
    def test_manyargs(self):
        # Test inspect.getsource with a regular function where the arguments
        # are on two lines and _not_ indented and the body on the second line
        # with the last arguments.
        self.assertEqual(inspect.getsource(mod.manyargs),
                         sourcerange(self.source, 17, 18))
    
    def test_twolinefunc(self):
        # Test inspect.getsource with a regular function where the body is on
        # two lines, following the argument list and continued on the next line
        # by a \\.
        self.assertEqual(inspect.getsource(mod.twolinefunc),
                         sourcerange(self.source, 21, 22))
    
    def test_lambda_in_list(self):
        # Test inspect.getsource with a one-line lambda function defined in a
        # list, indented.
        self.assertEqual(inspect.getsource(mod.a[1]),
                         sourcerange(self.source, 26, 26))
    
    def test_anonymous(self):
        # Test inspect.getsource with a lambda function defined as argument to
        # another function.
        self.assertEqual(inspect.getsource(mod.anonymous),
                         sourcerange(self.source, 32, 32))

def test_main():
    global mod
    file = open(TESTFN, "w")
    file.write(onelinesource)
    file.close()
    mod = imp.load_source("onelinemod", TESTFN)

    run_unittest(TestOneliners)
    
    for fname in [TESTFN, TESTFN + 'c', TESTFN + 'o']:
        try:
            os.unlink(fname)
        except OSError:
            pass

if __name__ == "__main__":
    test_main()