diff --git a/conftest.py b/conftest.py
index c20bde40916e7f7f5163c20e1c48c7b0f3d3b3a2..3bbdefe7c1f3e03ad87d07449d7d7214e18bfeba 100644
--- a/conftest.py
+++ b/conftest.py
@@ -1,8 +1,8 @@
 import os
 import runpy
 import sys
-import warnings
 import tempfile
+import warnings
 
 import nbformat
 import pytest
@@ -34,7 +34,7 @@ def add_path_to_ignore(path):
 
 
 collect_ignore = [os.path.join(SCRIPT_FOLDER, "doc", "conf.py"),
-        os.path.join(SCRIPT_FOLDER, "pystencils", "opencl", "opencl.autoinit")]
+                  os.path.join(SCRIPT_FOLDER, "pystencils", "opencl", "opencl.autoinit")]
 add_path_to_ignore('pystencils_tests/benchmark')
 add_path_to_ignore('_local_tmp')
 
@@ -83,8 +83,6 @@ for root, sub_dirs, files in os.walk('.'):
             collect_ignore.append(f)
 
 
-
-
 class IPythonMockup:
     def run_line_magic(self, *args, **kwargs):
         pass
diff --git a/pystencils/kernel_wrapper.py b/pystencils/kernel_wrapper.py
index 0e327711e5a355219cc2664ac9a6c8a02d88bc09..3494b52a9fd060bbccbfe493f165dcc7d63c8c06 100644
--- a/pystencils/kernel_wrapper.py
+++ b/pystencils/kernel_wrapper.py
@@ -1,11 +1,14 @@
-"""
-Light-weight wrapper around a compiled kernel
-"""
 import pystencils
 
 
 class KernelWrapper:
-    def __init__(self, kernel, parameters, ast_node):
+    """
+    Light-weight wrapper around a compiled kernel.
+
+    Can be called while still providing access to underlying AST.
+    """
+
+    def __init__(self, kernel, parameters, ast_node: pystencils.astnodes.KernelFunction):
         self.kernel = kernel
         self.parameters = parameters
         self.ast = ast_node