Fix and improve shortcode_atts extension (#265)

* Fix and improve shortcode_atts extension

* Update shortcode_atts.php

Author: IanDelMar

src/ShortcodeAttsDynamicFunctionReturnTypeExtension.php

@@ -11,7 +11,10 @@
 use PhpParser\Node\Expr\FuncCall;
 use PHPStan\Analyser\Scope;
 use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\Accessory\NonEmptyArrayType;
+use PHPStan\Type\ArrayType;
 use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
 use PHPStan\Type\StringType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
@@ -31,20 +34,49 @@ public function isFunctionSupported(FunctionReflection $functionReflection): boo
     public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): ?Type
     {
         $args = $functionCall->getArgs();
-        if ($args === []) {
+
+        if (count($args) < 2) {
             return null;
         }
 
-        $type = $scope->getType($args[0]->value);
+        $pairsType = $scope->getType($args[0]->value);
+        $attsType = $scope->getType($args[1]->value);
+
+        if ($attsType->isIterableAtLeastOnce()->no()) {
+            return $pairsType;
+        }
+
+        if (! $this->hasConstantArrays($pairsType)) {
+            return $this->resolveTypeForNonConstantPairs($pairsType);
+        }
 
-        if (count($type->getConstantArrays()) === 0) {
-            return $type;
+        if (! $this->hasConstantArrays($attsType)) {
+            return $this->resolveTypeForNonConstantAtts($pairsType);
         }
 
-        $returnType = [];
-        foreach ($type->getConstantArrays() as $constantArray) {
-            // shortcode_atts values are coming from the defined defaults or from the actual string shortcode attributes
-            $returnType[] = new ConstantArrayType(
+        return $this->resolveTypeForConstantAtts($pairsType, $attsType);
+    }
+
+    protected function resolveTypeForNonConstantPairs(Type $pairsType): Type
+    {
+        $keyType = $pairsType->getIterableKeyType();
+        $valueType = TypeCombinator::union(
+            $pairsType->getIterableValueType(),
+            new StringType()
+        );
+        $arrayType = new ArrayType($keyType, $valueType);
+
+        return $pairsType->isIterableAtLeastOnce()->yes()
+            ? TypeCombinator::intersect($arrayType, new NonEmptyArrayType())
+            : $arrayType;
+    }
+
+    protected function resolveTypeForNonConstantAtts(Type $pairsType): Type
+    {
+        $types = [];
+
+        foreach ($pairsType->getConstantArrays() as $constantArray) {
+            $types[] = new ConstantArrayType(
                 $constantArray->getKeyTypes(),
                 array_map(
                     static function (Type $valueType): Type {
@@ -55,6 +87,52 @@ static function (Type $valueType): Type {
             );
         }
 
-        return TypeCombinator::union(...$returnType);
+        return TypeCombinator::union(...$types);
+    }
+
+    protected function resolveTypeForConstantAtts(Type $pairsType, Type $attsType): Type
+    {
+        $types = [];
+
+        foreach ($pairsType->getConstantArrays() as $constantPairsArray) {
+            foreach ($attsType->getConstantArrays() as $constantAttsArray) {
+                $types[] = $this->mergeArrays($constantPairsArray, $constantAttsArray);
+            }
+        }
+
+        return TypeCombinator::union(...$types);
+    }
+
+    protected function mergeArrays(ConstantArrayType $pairsArray, ConstantArrayType $attsArray): Type
+    {
+        if (count($attsArray->getKeyTypes()) === 0) {
+            return $pairsArray;
+        }
+
+        $builder = ConstantArrayTypeBuilder::createFromConstantArray($pairsArray);
+
+        foreach ($pairsArray->getKeyTypes() as $keyType) {
+            $hasOffsetValueType = $attsArray->hasOffsetValueType($keyType);
+
+            if ($hasOffsetValueType->no()) {
+                continue;
+            }
+
+            $valueType = $hasOffsetValueType->yes()
+                ? $attsArray->getOffsetValueType($keyType)
+                : TypeCombinator::union(
+                    $pairsArray->getOffsetValueType($keyType),
+                    $attsArray->getOffsetValueType($keyType)
+                );
+
+            $builder->setOffsetValueType($keyType, $valueType);
+        }
+
+        return $builder->getArray();
+    }
+
+    protected function hasConstantArrays(Type $type): bool
+    {
+        return count($type->getConstantArrays()) > 0;
     }
 }

tests/data/shortcode_atts.php

@@ -4,14 +4,41 @@
 
 namespace SzepeViktor\PHPStan\WordPress\Tests;
 
+use function shortcode_atts;
 use function PHPStan\Testing\assertType;
 
-$atts = $_GET['atts'] ?? [];
+/** @var array $pairs */
+$pairs = $_GET['pairs'];
+
+/** @var array $atts */
+$atts = $_GET['atts'];
 
 // shortcode_atts filters atts by returning all atts that are occurring in the pairs
 // As atts are supposed to be strings the function is expected to return the pair type or a string
 
+// Constant $pairs and non-constant $atts
 assertType('array{}', shortcode_atts([], $atts));
 assertType('array{foo: string, bar: string}', shortcode_atts(['foo' => '', 'bar' => ''], $atts));
 assertType('array{foo: string, bar: 19|string}', shortcode_atts(['foo' => 'foo-value', 'bar' => 19], $atts));
 assertType('array{foo: string|null, bar: 17|string, baz: string}', shortcode_atts(['foo' => null, 'bar' => 17, 'baz' => ''], $atts));
+
+// Constant $pairs and constant $atts
+assertType("array{foo: '', bar: '19', baz: 'aString'}", shortcode_atts(['foo' => null, 'bar' => 17, 'baz' => ''], ['foo' => '', 'bar' => '19', 'baz' => 'aString']));
+
+// Constant $pairs and empty $atts
+assertType('array{}', shortcode_atts([], []));
+assertType("array{foo: '', bar: ''}", shortcode_atts(['foo' => '', 'bar' => ''], []));
+assertType("array{foo: 'foo-value', bar: 19}", shortcode_atts(['foo' => 'foo-value', 'bar' => 19], []));
+assertType("array{foo: null, bar: 17, baz: ''}", shortcode_atts(['foo' => null, 'bar' => 17, 'baz' => ''], []));
+
+// Non-constant $pairs (mixed) and varying $atts
+assertType('array', shortcode_atts($pairs, $atts));
+assertType('array', shortcode_atts($pairs, []));
+assertType('array', shortcode_atts($pairs, ['foo' => '', 'bar' => '19', 'baz' => '']));
+
+// Non-constant $pairs (int) and varying $atts
+/** @var array<string, int> $pairs */
+$pairs = $_GET['pairs'];
+assertType('array<string, int|string>', shortcode_atts($pairs, $atts));
+assertType('array<string, int>', shortcode_atts($pairs, []));
+assertType('array<string, int|string>', shortcode_atts($pairs, ['foo' => '', 'bar' => '19', 'baz' => '']));