diff --git a/integration/CMakeDemo/kernels.py b/integration/CMakeDemo/kernels.py
index 5fcac7c8ea3c1f71e185f6f5f70af9b16c948096..517c7e0f512761c2cd08842ca5d1a9cd1924edba 100644
--- a/integration/CMakeDemo/kernels.py
+++ b/integration/CMakeDemo/kernels.py
@@ -4,10 +4,12 @@ import sympy as sp
 
 from pystencils import fields, kernel
 
-from pystencilssfg import SourceFileGenerator
+from pystencilssfg import SourceFileGenerator, SfgComposer
 
 
-with SourceFileGenerator() as sfg:
+with SourceFileGenerator() as ctx:
+    sfg = SfgComposer(ctx)
+
     src, dst = fields("src, dst(1) : double[2D]")
 
     @kernel
diff --git a/integration/CMakeDemo/more_kernels.py b/integration/CMakeDemo/more_kernels.py
index 75b798fb7b30431613531edee033df1512d23212..5fbd6addccb66596c289ef9abc38b7baf772226b 100644
--- a/integration/CMakeDemo/more_kernels.py
+++ b/integration/CMakeDemo/more_kernels.py
@@ -4,10 +4,11 @@ import sympy as sp
 
 from pystencils import fields, kernel, Field
 
-from pystencilssfg import SourceFileGenerator
+from pystencilssfg import SourceFileGenerator, SfgComposer
 
+with SourceFileGenerator() as ctx:
+    sfg = SfgComposer(ctx)
 
-with SourceFileGenerator() as sfg:
     src: Field = fields("src: double[2D]")
 
     h = sp.Symbol('h')
diff --git a/integration/MakeDemo/kernels.py b/integration/MakeDemo/kernels.py
index 310b9bcd0c244afe86c869117b19c4b335ef6c34..43c0f858cdd6b9bf58566fcad342713d48e81394 100644
--- a/integration/MakeDemo/kernels.py
+++ b/integration/MakeDemo/kernels.py
@@ -4,14 +4,16 @@ import sympy as sp
 
 from pystencils import fields, kernel
 
-from pystencilssfg import SourceFileGenerator, SfgConfiguration
+from pystencilssfg import SourceFileGenerator, SfgConfiguration, SfgComposer
 from pystencilssfg.source_concepts.cpp import mdspan_ref
 
 sfg_config = SfgConfiguration(
     outer_namespace="make_demo"
 )
 
-with SourceFileGenerator(sfg_config) as sfg:
+with SourceFileGenerator(sfg_config) as ctx:
+    sfg = SfgComposer(ctx)
+    
     sfg.prelude("""Generated by the pystencils Source File Generator.
 
 Author: Frederik Hennig <frederik.hennig@fau.de>""")
diff --git a/integration/TestSequencing.py b/integration/TestSequencing.py
index 9e97052c5925e5507827906a7b206e3f87b8575f..3a0e040bfb0be6d97c1e5896ee0a201c905ec466 100644
--- a/integration/TestSequencing.py
+++ b/integration/TestSequencing.py
@@ -1,9 +1,10 @@
-from pystencilssfg import SourceFileGenerator
+from pystencilssfg import SourceFileGenerator, SfgComposer
 
 from lbmpy.advanced_streaming import Timestep
 from lbmpy import LBMConfig, create_lb_ast
 
-with SourceFileGenerator() as sfg:
+with SourceFileGenerator() as ctx:
+    sfg = SfgComposer(ctx)
 
     lb_config = LBMConfig(streaming_pattern='esotwist')
 
diff --git a/integration/test_classes.py b/integration/test_classes.py
index 40d7098fd848458ba8f4e162eec64f6589ab4cfd..ecd6b89510485827ae4b0a8ffcbe1f0ccf9f634f 100644
--- a/integration/test_classes.py
+++ b/integration/test_classes.py
@@ -1,5 +1,5 @@
 # type: ignore
-from pystencilssfg import SourceFileGenerator, SfgConfiguration
+from pystencilssfg import SourceFileGenerator, SfgConfiguration, SfgComposer
 from pystencilssfg.configuration import SfgCodeStyle
 from pystencilssfg.source_concepts import SrcObject
 from pystencilssfg.source_components import SfgClass, SfgMemberVariable, SfgConstructor, SfgMethod, SfgVisibility
@@ -17,7 +17,8 @@ sfg_config = SfgConfiguration(
 
 f, g = fields("f, g(1): double[2D]")
 
-with SourceFileGenerator(sfg_config) as sfg:
+with SourceFileGenerator(sfg_config) as ctx:
+    sfg = SfgComposer(ctx)
 
     @kernel
     def assignments():
diff --git a/src/pystencilssfg/generator.py b/src/pystencilssfg/generator.py
index cea341b5570efa0954b1a6bc8c41f07e4a652049..6a2038c289584d54585437bb1b025eebdede434a 100644
--- a/src/pystencilssfg/generator.py
+++ b/src/pystencilssfg/generator.py
@@ -7,7 +7,6 @@ from os import path
 
 from .configuration import SfgConfiguration, config_from_commandline, merge_configurations
 from .context import SfgContext
-from .composer import SfgComposer
 
 
 class SourceFileGenerator:
@@ -35,9 +34,9 @@ class SourceFileGenerator:
             if path.exists(file):
                 os.remove(file)
 
-    def __enter__(self):
+    def __enter__(self) -> SfgContext:
         self.clean_files()
-        return SfgComposer(self._context)
+        return self._context
 
     def __exit__(self, exc_type, exc_value, traceback):
         if exc_type is None: