diff --git a/.gitignore b/.gitignore
index 1716623..a99044e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,7 @@
 
 # Documents folder
 doc
-.yardoc
\ No newline at end of file
+.yardoc
+
+# Ignore local Bundler installations
+/.bundle/
diff --git a/Gemfile.lock b/Gemfile.lock
index 7b27356..75d3b32 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    yard-cucumber (3.1.0)
+    yard-cucumber (4.0.0)
       cucumber (>= 2.0, < 4.0)
       gherkin (>= 4.0, < 6.0)
       yard (~> 0.8, >= 0.8.1)
diff --git a/lib/templates/default/fulldoc/html/setup.rb b/lib/templates/default/fulldoc/html/setup.rb
index b38f871..5a6c236 100644
--- a/lib/templates/default/fulldoc/html/setup.rb
+++ b/lib/templates/default/fulldoc/html/setup.rb
@@ -163,7 +163,7 @@ def generate_full_list(objects,type,options = {})
 # This method removes the namespace from the root node, generates the class list,
 # and then adds it back into the root node.
 #
-def class_list(root = Registry.root, tree = TreeContext.new)
+def class_list(root = Registry.root, tree = self.class.const_get(:TreeContext).new)
   return super unless root == Registry.root
 
   cucumber_namespace = YARD::CodeObjects::Cucumber::CUCUMBER_NAMESPACE