Fix a minor issue with inline import statementsv2.1.1rc2

This will hopefully be the last issue related to these changes, added two more unit tests to cover these cases
and ran the whole suite successfully.

This may not have actually been the issue reported by mobiledreamers@gmail.com, but investigating his issue led
to these changes

Signed-off-by: R. Tyler Ballance <tyler@slide.com>

2 files changed, 35 insertions, 9 deletions

diff --git a/src/Compiler.py b/src/Compiler.py
index 9d4023f..a885d07 100644
--- a/src/Compiler.py
+++ b/src/Compiler.py
@@ -1843,7 +1843,10 @@ class ModuleCompiler(SettingsManager, GenUtils):
         self._specialVars[name] = contents.strip()
 
     def addImportStatement(self, impStatement):
-        self._importStatements.append(impStatement)
+        if not self._methodBodyChunks:
+            # In the case where we are importing inline in the middle of a source block
+            # we don't want to inadvertantly import the module at the top of the file either
+            self._importStatements.append(impStatement)
 
         #@@TR 2005-01-01: there's almost certainly a cleaner way to do this!
         importVarNames = impStatement[impStatement.find('import') + len('import'):].split(',')
diff --git a/src/Tests/Regressions.py b/src/Tests/Regressions.py
index fb8327f..af49d41 100644
--- a/src/Tests/Regressions.py
+++ b/src/Tests/Regressions.py
@@ -41,14 +41,14 @@ class GetAttrTest(unittest.TestCase):
         self.failUnlessRaises(GetAttrException, template.raiseme)
 
 
-class InlineFromImportTest(unittest.TestCase):
-    '''
-        Verify that a bug introduced in v2.1.0 where an inline:
-            #from module import class
-        would result in the following code being generated:
-            improt class
-    '''
-    def runTest(self):
+class InlineImportTest(unittest.TestCase):
+    def test_FromFooImportThing(self):
+        '''
+            Verify that a bug introduced in v2.1.0 where an inline:
+                #from module import class
+            would result in the following code being generated:
+                import class
+        '''
         template = '''
             #def myfunction()
                 #if True
@@ -66,6 +66,29 @@ class InlineFromImportTest(unittest.TestCase):
         rc = template.myfunction()
         assert rc == 17, (template, 'Didn\'t get a proper return value')
 
+    def test_ImportFailModule(self):
+        template = '''
+            #try
+                #import invalidmodule
+            #except
+                #set invalidmodule = dict(FOO='BAR!')
+            #end try
+
+            $invalidmodule.FOO
+        '''
+        template = Cheetah.Template.Template.compile(template, compilerSettings={}, keepRefToGeneratedCode=True)
+        template = template(searchList=[{}])
+
+        assert template, 'We should have a valid template object by now'
+        assert str(template), 'We weren\'t able to properly generate the result from the template'
+
+    def test_ProperImportOfBadModule(self):
+        template = '''
+            #from invalid import fail
+                
+            This should totally $fail
+        '''
+        self.failUnlessRaises(ImportError, Cheetah.Template.Template.compile, template, compilerSettings={}, keepRefToGeneratedCode=True)
 
 if __name__ == '__main__':
     unittest.main()