diff --git a/extension.neon b/extension.neon
index 6a34985..8bea503 100644
--- a/extension.neon
+++ b/extension.neon
@@ -11,7 +11,6 @@ parameters:
 		- stubs/Application/UI/Presenter.stub
 		- stubs/Caching/Cache.stub
 		- stubs/ComponentModel/Component.stub
-		- stubs/ComponentModel/Container.stub
 		- stubs/ComponentModel/IComponent.stub
 		- stubs/ComponentModel/IContainer.stub
 		- stubs/Database/ResultSet.stub
@@ -127,3 +126,8 @@ services:
 		class: PHPStan\Type\Nette\StringsReplaceCallbackClosureTypeExtension
 		tags:
 			- phpstan.staticMethodParameterClosureTypeExtension
+
+	-
+		class: PHPStan\Stubs\Nette\StubFilesExtensionLoader
+		tags:
+			- phpstan.stubFilesExtension
diff --git a/src/Stubs/Nette/StubFilesExtensionLoader.php b/src/Stubs/Nette/StubFilesExtensionLoader.php
new file mode 100644
index 0000000..342ffb0
--- /dev/null
+++ b/src/Stubs/Nette/StubFilesExtensionLoader.php
@@ -0,0 +1,47 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Stubs\Nette;
+
+use Composer\InstalledVersions;
+use OutOfBoundsException;
+use PHPStan\PhpDoc\StubFilesExtension;
+use function class_exists;
+use function dirname;
+use function version_compare;
+
+class StubFilesExtensionLoader implements StubFilesExtension
+{
+
+	public function getFiles(): array
+	{
+		$stubsDir = dirname(dirname(dirname(__DIR__))) . '/stubs';
+		$path = $stubsDir;
+
+		$files = [];
+
+		$componentModelVersion = self::getInstalledVersion('nette/component-model');
+		if ($componentModelVersion !== null && version_compare($componentModelVersion, '3.1.0', '>=')) {
+			$files[] = $path . '/ComponentModel/Container_3_1.stub';
+		} else {
+			$files[] = $path . '/ComponentModel/Container.stub';
+		}
+
+		return $files;
+	}
+
+	private static function getInstalledVersion(string $package): ?string
+	{
+		if (!class_exists(InstalledVersions::class)) {
+			return null;
+		}
+
+		try {
+			$installedVersion = InstalledVersions::getVersion($package);
+		} catch (OutOfBoundsException $e) {
+			return null;
+		}
+
+		return $installedVersion;
+	}
+
+}
diff --git a/stubs/ComponentModel/Container.stub b/stubs/ComponentModel/Container.stub
index 2c0ba4b..384b2b6 100644
--- a/stubs/ComponentModel/Container.stub
+++ b/stubs/ComponentModel/Container.stub
@@ -10,7 +10,7 @@ class Container extends Component implements IContainer
 	 * @phpstan-param null|class-string<T> $filterType
 	 * @phpstan-return ($filterType is null ? \Iterator<int|string, \Nette\ComponentModel\IComponent> : \Iterator<int|string, T>)
 	 */
-	public function getComponents(bool $deep = false, string $filterType = null): \Iterator
+	public function getComponents(bool $deep = false, ?string $filterType = null): \Iterator
 	{
 		// nothing
 	}
diff --git a/stubs/ComponentModel/Container_3_1.stub b/stubs/ComponentModel/Container_3_1.stub
new file mode 100644
index 0000000..18d16eb
--- /dev/null
+++ b/stubs/ComponentModel/Container_3_1.stub
@@ -0,0 +1,21 @@
+<?php
+
+namespace Nette\ComponentModel;
+
+class Container extends Component implements IContainer
+{
+
+	/**
+	 * @template T of \Nette\ComponentModel\IComponent
+	 * @phpstan-param null|class-string<T> $filterType
+	 * @phpstan-return (
+	 * 	$deep is false
+	 * 	? ($filterType is null ? array<int|string, \Nette\ComponentModel\IComponent> : array<int|string, T>)
+	 * 	: ($filterType is null ? \Iterator<int|string, \Nette\ComponentModel\IComponent> : \Iterator<int|string, T>)
+	 * )
+	 */
+	public function getComponents(bool $deep = false, ?string $filterType = null): iterable
+	{
+		// nothing
+	}
+}
diff --git a/tests/Type/Nette/ComponentModelContainerDynamicReturnTypeExtensionTest.php b/tests/Type/Nette/ComponentModelContainerDynamicReturnTypeExtensionTest.php
new file mode 100644
index 0000000..c73f3e8
--- /dev/null
+++ b/tests/Type/Nette/ComponentModelContainerDynamicReturnTypeExtensionTest.php
@@ -0,0 +1,61 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Nette;
+
+use Composer\InstalledVersions;
+use OutOfBoundsException;
+use PHPStan\Testing\TypeInferenceTestCase;
+use function class_exists;
+use function version_compare;
+
+class ComponentModelContainerDynamicReturnTypeExtensionTest extends TypeInferenceTestCase
+{
+
+	/**
+	 * @return iterable<string, mixed[]>
+	 */
+	public function dataFileAsserts(): iterable
+	{
+		$componentModelVersion = self::getInstalledVersion('nette/component-model');
+
+		$suffix = $componentModelVersion !== null && version_compare($componentModelVersion, '3.1.0', '>=') ? '31' : '';
+
+		yield from self::gatherAssertTypes(__DIR__ . '/data/componentModelContainer' . $suffix . '.php');
+	}
+
+	/**
+	 * @dataProvider dataFileAsserts
+	 * @param mixed ...$args
+	 */
+	public function testFileAsserts(
+		string $assertType,
+		string $file,
+		...$args
+	): void
+	{
+		$this->assertFileAsserts($assertType, $file, ...$args);
+	}
+
+	public static function getAdditionalConfigFiles(): array
+	{
+		return [
+			__DIR__ . '/phpstan.neon',
+		];
+	}
+
+	private static function getInstalledVersion(string $package): ?string
+	{
+		if (!class_exists(InstalledVersions::class)) {
+			return null;
+		}
+
+		try {
+			$installedVersion = InstalledVersions::getVersion($package);
+		} catch (OutOfBoundsException $e) {
+			return null;
+		}
+
+		return $installedVersion;
+	}
+
+}
diff --git a/tests/Type/Nette/data/componentModelContainer.php b/tests/Type/Nette/data/componentModelContainer.php
new file mode 100644
index 0000000..f082254
--- /dev/null
+++ b/tests/Type/Nette/data/componentModelContainer.php
@@ -0,0 +1,17 @@
+<?php
+
+namespace PHPStan\Type\Nette\Data\ComponentModel;
+
+use Nette\Application\UI\Form;
+use Nette\Forms\Container;
+use function PHPStan\Testing\assertType;
+
+class SomeForm extends Form {
+}
+
+$someForm = new SomeForm();
+
+assertType('Iterator<int|string, Nette\ComponentModel\IComponent>', $someForm->getComponents(false));
+assertType('Iterator<int|string, Nette\Forms\Container>', $someForm->getComponents(false, Container::class));
+assertType('Iterator<int|string, Nette\ComponentModel\IComponent>', $someForm->getComponents(true));
+assertType('Iterator<int|string, Nette\Forms\Container>', $someForm->getComponents(true, Container::class));
diff --git a/tests/Type/Nette/data/componentModelContainer31.php b/tests/Type/Nette/data/componentModelContainer31.php
new file mode 100644
index 0000000..a2240d1
--- /dev/null
+++ b/tests/Type/Nette/data/componentModelContainer31.php
@@ -0,0 +1,17 @@
+<?php
+
+namespace PHPStan\Type\Nette\Data\ComponentModel;
+
+use Nette\Application\UI\Form;
+use Nette\Forms\Container;
+use function PHPStan\Testing\assertType;
+
+class SomeForm extends Form {
+}
+
+$someForm = new SomeForm();
+
+assertType('array<int|string, Nette\ComponentModel\IComponent>', $someForm->getComponents(false));
+assertType('array<int|string, Nette\Forms\Container>', $someForm->getComponents(false, Container::class));
+assertType('Iterator<int|string, Nette\ComponentModel\IComponent>', $someForm->getComponents(true));
+assertType('Iterator<int|string, Nette\Forms\Container>', $someForm->getComponents(true, Container::class));