Add null to the container get method return type when ContainerInterace::NULL_ON_INVALID_REFERENCE is passed as the second argument. (#610)

Add null to the container get method return type when ContainerInterface::NULL_ON_INVALID_REFERENCE is passed as the second argument.

Author: mad-briller

src/Type/ContainerDynamicReturnTypeExtension.php

@@ -9,9 +9,10 @@
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\Type\Constant\ConstantBooleanType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\NullType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
-use Psr\Container\ContainerInterface;
+use Symfony\Component\DependencyInjection\ContainerInterface;
 
 class ContainerDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -41,24 +42,54 @@ public function getTypeFromMethodCall(
         Scope $scope
     ): Type {
         $returnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
-        $args = $methodCall->getArgs();
-        if (count($args) !== 1) {
-            return $returnType;
-        }
-
         $methodName = $methodReflection->getName();
-        $types = [];
-        $argType = $scope->getType($args[0]->value);
 
-        foreach ($argType->getConstantStrings() as $constantStringType) {
-            $serviceId = $constantStringType->getValue();
-            $service = $this->serviceMap->getService($serviceId);
-            if ($methodName === 'get') {
-                $types[] = $service !== null ? $service->getType() : $returnType;
-            } elseif ($methodName === 'has') {
+        if ($methodName === 'has') {
+            $args = $methodCall->getArgs();
+            if (count($args) !== 1) {
+                return $returnType;
+            }
+
+            $types = [];
+            $argType = $scope->getType($args[0]->value);
+
+            foreach ($argType->getConstantStrings() as $constantStringType) {
+                $serviceId = $constantStringType->getValue();
+                $service = $this->serviceMap->getService($serviceId);
                 $types[] = new ConstantBooleanType($service !== null);
             }
+
+            return TypeCombinator::union(...$types);
+        } elseif ($methodName === 'get') {
+            $args = $methodCall->getArgs();
+            if (count($args) === 0) {
+                return $returnType;
+            }
+
+            $types = [];
+
+            if (isset($args[1])) {
+                $invalidBehaviour = $scope->getType($args[1]->value);
+
+                foreach ($invalidBehaviour->getConstantScalarValues() as $value) {
+                    if ($value === ContainerInterface::NULL_ON_INVALID_REFERENCE) {
+                        $types[] = new NullType();
+                        break;
+                    }
+                }
+            }
+
+            $argType = $scope->getType($args[0]->value);
+
+            foreach ($argType->getConstantStrings() as $constantStringType) {
+                $serviceId = $constantStringType->getValue();
+                $service = $this->serviceMap->getService($serviceId);
+                $types[] = $service !== null ? $service->getType() : $returnType;
+            }
+
+            return TypeCombinator::union(...$types);
         }
-        return TypeCombinator::union(...$types);
+
+        return $returnType;
     }
 }

tests/src/Type/DrupalContainerDynamicReturnTypeTest.php

@@ -17,6 +17,7 @@ public function dataFileAsserts(): iterable
         yield from self::gatherAssertTypes(__DIR__ . '/data/drupal-service-static.php');
         yield from self::gatherAssertTypes(__DIR__ . '/data/drupal-class-resolver.php');
         yield from self::gatherAssertTypes(__DIR__ . '/data/bug-563.php');
+        yield from self::gatherAssertTypes(__DIR__ . '/data/container-optional.php');
     }
 
     /**

tests/src/Type/data/container-optional.php

@@ -0,0 +1,41 @@
+<?php
+
+namespace DrupalContainerOptional;
+
+use Symfony\Component\DependencyInjection\ContainerInterface;
+use Drupal\service_map\MyService;
+use function PHPStan\Testing\assertType;
+
+function test(): void {
+    $container = \Drupal::getContainer();
+
+    assertType(
+        'object|null',
+        $container->get('unknown_service', ContainerInterface::NULL_ON_INVALID_REFERENCE)
+    );
+
+    assertType(
+        MyService::class . '|null',
+        $container->get('service_map.my_service', ContainerInterface::NULL_ON_INVALID_REFERENCE)
+    );
+
+    assertType(
+        MyService::class,
+        $container->get('service_map.my_service', ContainerInterface::EXCEPTION_ON_INVALID_REFERENCE)
+    );
+
+    assertType(
+        MyService::class,
+        $container->get('service_map.my_service', ContainerInterface::RUNTIME_EXCEPTION_ON_INVALID_REFERENCE)
+    );
+
+    assertType(
+        MyService::class,
+        $container->get('service_map.my_service', ContainerInterface::IGNORE_ON_INVALID_REFERENCE)
+    );
+
+    assertType(
+        MyService::class,
+        $container->get('service_map.my_service', ContainerInterface::IGNORE_ON_UNINITIALIZED_REFERENCE)
+    );
+}