diff --git a/composer.json b/composer.json
index dbc32e5..cf9b51d 100644
--- a/composer.json
+++ b/composer.json
@@ -12,7 +12,8 @@
   ],
   "require": {
     "ocramius/proxy-manager": "^2.0",
-    "nette/di": "^2.4"
+    "nette/di": "^2.4",
+    "nette/php-generator": "^2.6|^3.0"
   },
   "require-dev": {
     "phpunit/phpunit": "^5.6",
diff --git a/src/DI/ProxyExtension.php b/src/DI/ProxyExtension.php
index 38ba69f..0822fee 100644
--- a/src/DI/ProxyExtension.php
+++ b/src/DI/ProxyExtension.php
@@ -9,7 +9,7 @@
 use Nette\DI\Helpers;
 use Nette\DI\Statement;
 use Nette\PhpGenerator\ClassType;
-use Nette\PhpGenerator\Method;
+use Nette\PhpGenerator\Closure;
 use Nette\Utils\Validators;
 use ProxyManager\Configuration;
 use ProxyManager\Factory\LazyLoadingValueHolderFactory;
@@ -116,8 +116,8 @@ public function afterCompile(ClassType $class)
 			// modify original method body to return proxy instead
 			$method = $class->getMethod(Container::getMethodName($name));
 			$method->setBody(sprintf(
-				"return \$this->getService('%s')->createProxy(\n\t%s::class,\n\tfunction (&\$wrappedObject, \$proxy, \$method, \$parameters, &\$initializer) {\n\t\t\$wrappedObject = (%s)();\n\t\t\$initializer = null;\n\t}\n);",
-				$this->prefix('lazyLoadingValueHolderFactory'), $type, ltrim(preg_replace('#^#m', "\t\t", (new Method())->addBody($method->getBody())))
+				"return \$this->getService('%s')->createProxy(\n\t%s::class,\n\tfunction (&\$wrappedObject, \$proxy, \$method, \$parameters, &\$initializer) {\n\t\t\$wrappedObject = (%s)();\n\t\t\$initializer = null;\n\t\treturn true;\n\t}\n);",
+				$this->prefix('lazyLoadingValueHolderFactory'), $type, ltrim(preg_replace('#^#m', "\t\t", (new Closure())->addBody($method->getBody())))
 			));
 		}