Define psalm types for most pipeline stages

doctrine · Mar 23, 2023 · b7deae2 · b7deae2
1 parent 77d5827
commit b7deae2
Show file tree

Hide file tree

Showing 31 changed files with 1,174 additions and 347 deletions.
diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Aggregation.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Aggregation.php
@@ -17,6 +17,7 @@
 use function array_merge;
 use function assert;
 
+/** @psalm-import-type PipelineExpression from Builder */
 final class Aggregation implements IteratorAggregate
 {
     private DocumentManager $dm;
@@ -25,7 +26,10 @@ final class Aggregation implements IteratorAggregate
 
     private Collection $collection;
 
-    /** @var array<string, mixed> */
+    /**
+     * @var array<string, mixed>
+     * @psalm-var PipelineExpression
+     */
     private array $pipeline;
 
     /** @var array<string, mixed> */
@@ -36,6 +40,7 @@ final class Aggregation implements IteratorAggregate
     /**
      * @param array<string, mixed> $pipeline
      * @param array<string, mixed> $options
+     * @psalm-param PipelineExpression $pipeline
      */
     public function __construct(DocumentManager $dm, ?ClassMetadata $classMetadata, Collection $collection, array $pipeline, array $options = [], bool $rewindable = true)
     {

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Builder.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Builder.php
@@ -29,6 +29,8 @@
  * Fluent interface for building aggregation pipelines.
  *
  * @psalm-import-type SortShape from Sort
+ * @psalm-import-type StageExpression from Stage
+ * @psalm-type PipelineExpression = list<StageExpression>
  */
 class Builder
 {
@@ -269,6 +271,7 @@ public function getAggregation(array $options = []): Aggregation
      * given.
      *
      * @return array<array<string, mixed>>
+     * @psalm-return PipelineExpression
      */
     // phpcs:enable Squiz.Commenting.FunctionComment.ExtraParamComment
     public function getPipeline(/* bool $applyFilters = true */): array

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Expr.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Expr.php
@@ -24,6 +24,8 @@
 
 /**
  * Fluent interface for building aggregation pipelines.
+ *
+ * @psalm-type OperatorExpression = array<string, mixed>|object
  */
 class Expr implements GenericOperatorsInterface
 {

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage.php
@@ -13,6 +13,9 @@
  * Fluent interface for building aggregation pipelines.
  *
  * @internal
+ *
+ * @psalm-import-type PipelineExpression from Builder
+ * @psalm-type StageExpression = non-empty-array<non-empty-string, mixed>|object
  */
 abstract class Stage
 {
@@ -186,6 +189,7 @@ public function geoNear($x, $y = null): Stage\GeoNear
      * Returns the assembled aggregation pipeline
      *
      * @return array<array<string, mixed>>
+     * @psalm-return PipelineExpression
      */
     public function getPipeline(): array
     {

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/AbstractReplace.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/AbstractReplace.php
@@ -0,0 +1,64 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Doctrine\ODM\MongoDB\Aggregation\Stage;
+
+use Doctrine\ODM\MongoDB\Aggregation\Builder;
+use Doctrine\ODM\MongoDB\Aggregation\Expr;
+use Doctrine\ODM\MongoDB\DocumentManager;
+use Doctrine\ODM\MongoDB\Mapping\ClassMetadata;
+use Doctrine\ODM\MongoDB\Persisters\DocumentPersister;
+use Doctrine\ODM\MongoDB\Types\Type;
+
+use function array_map;
+use function is_array;
+use function is_string;
+use function substr;
+
+abstract class AbstractReplace extends Operator
+{
+    /** @var string|mixed[]|Expr|null */
+    protected $expression;
+    protected DocumentManager $dm;
+    protected ClassMetadata $class;
+
+    /** @param string|mixed[]|Expr|null $expression */
+    final public function __construct(Builder $builder, DocumentManager $documentManager, ClassMetadata $class, $expression = null)
+    {
+        Operator::__construct($builder);
+
+        $this->dm         = $documentManager;
+        $this->class      = $class;
+        $this->expression = $expression;
+    }
+
+    private function getDocumentPersister(): DocumentPersister
+    {
+        return $this->dm->getUnitOfWork()->getDocumentPersister($this->class->name);
+    }
+
+    /** @return array|string */
+    protected function getReplaceExpression()
+    {
+        return $this->expression !== null ? $this->convertExpression($this->expression) : $this->expr->getExpression();
+    }
+
+    /**
+     * @param mixed[]|string|mixed $expression
+     *
+     * @return mixed[]|string|mixed
+     */
+    private function convertExpression($expression)
+    {
+        if (is_array($expression)) {
+            return array_map([$this, 'convertExpression'], $expression);
+        }
+
+        if (is_string($expression) && substr($expression, 0, 1) === '$') {
+            return '$' . $this->getDocumentPersister()->prepareFieldName(substr($expression, 1));
+        }
+
+        return Type::convertPHPToDatabaseValue(Expr::convertExpression($expression));
+    }
+}
diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/AddFields.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/AddFields.php
@@ -4,11 +4,17 @@
 
 namespace Doctrine\ODM\MongoDB\Aggregation\Stage;
 
+use Doctrine\ODM\MongoDB\Aggregation\Expr;
+
 /**
  * Fluent interface for adding a $addFields stage to an aggregation pipeline.
+ *
+ * @psalm-import-type OperatorExpression from Expr
+ * @psalm-type AddFieldsStageExpression = array{'$addFields': array<string, OperatorExpression|mixed>}
  */
 final class AddFields extends Operator
 {
+    /** @return AddFieldsStageExpression */
     public function getExpression(): array
     {
         return ['$addFields' => $this->expr->getExpression()];

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/CollStats.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/CollStats.php
@@ -9,6 +9,13 @@
 
 /**
  * Fluent interface for adding a $collStats stage to an aggregation pipeline.
+ *
+ * @psalm-type CollStatsStageExpression = array{
+ *     '$collStats': array{
+ *         latencyStats?: array{histograms?: bool},
+ *         storageStats?: array{},
+ *     }
+ * }
  */
 class CollStats extends Stage
 {
@@ -45,6 +52,7 @@ public function showStorageStats(): self
         return $this;
     }
 
+    /** @psalm-return CollStatsStageExpression */
     public function getExpression(): array
     {
         $collStats = [];

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Count.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Count.php
@@ -9,6 +9,8 @@
 
 /**
  * Fluent interface for adding a $count stage to an aggregation pipeline.
+ *
+ * @psalm-type CountStageExpression = array{'$count': string}
  */
 class Count extends Stage
 {
@@ -21,6 +23,7 @@ public function __construct(Builder $builder, string $fieldName)
         $this->fieldName = $fieldName;
     }
 
+    /** @psalm-return CountStageExpression */
     public function getExpression(): array
     {
         return [

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Densify.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Densify.php
@@ -12,6 +12,20 @@
 
 /**
  * Fluent interface for adding a $densify stage to an aggregation pipeline.
+ *
+ * @psalm-type BoundsType = 'full'|'partition'|array{0: int|float|UTCDateTime, 1: int|float|UTCDateTime}
+ * @psalm-type UnitType = 'year'|'month'|'week'|'day'|'hour'|'minute'|'second'|'millisecond'
+ * @psalm-type DensifyStageExpression = array{
+ *     '$densify': object{
+ *         field: string,
+ *         partitionByFields?: list<string>,
+ *         range?: object{
+ *             bounds: BoundsType,
+ *             step: int|float,
+ *             unit?: UnitType
+ *         }
+ *     }
+ * }
  */
 class Densify extends Stage
 {
@@ -37,8 +51,10 @@ public function partitionByFields(string ...$fields): self
     }
 
     /**
-     * @param array{0: int|float|UTCDateTime, 1: int|float|UTCDateTime}|string $bounds
-     * @param int|float                                                        $step
+     * @param array|string $bounds
+     * @param int|float    $step
+     * @psalm-param BoundsType  $bounds
+     * @psalm-param ''|UnitType $unit
      */
     public function range($bounds, $step, string $unit = ''): self
     {
@@ -54,6 +70,7 @@ public function range($bounds, $step, string $unit = ''): self
         return $this;
     }
 
+    /** @psalm-return DensifyStageExpression */
     public function getExpression(): array
     {
         $params = (object) ['field' => $this->field];

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Facet.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Facet.php
@@ -13,6 +13,9 @@
 
 /**
  * Fluent interface for adding a $facet stage to an aggregation pipeline.
+ *
+ * @psalm-import-type PipelineExpression from Builder
+ * @psalm-type FacetStageExpression = array{'$facet': array<string, PipelineExpression>}
  */
 class Facet extends Stage
 {

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Fill.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Fill.php
@@ -17,7 +17,18 @@
 /**
  * Fluent interface for adding a $fill stage to an aggregation pipeline.
  *
- * @psalm-type SortShape = array<string, int|string>
+ * @psalm-import-type SortDirectionKeywords from Sort
+ * @psalm-import-type OperatorExpression from Expr
+ * @psalm-type SortDirection = int|SortDirectionKeywords
+ * @psalm-type SortShape = array<string, SortDirection>
+ * @psalm-type FillStageExpression = array{
+ *     '$fill': array{
+ *         partitionBy?: string|OperatorExpression,
+ *         partitionByFields?: list<string>,
+ *         sortBy?: SortShape,
+ *         output?: array,
+ *     }
+ * }
  */
 class Fill extends Stage
 {
@@ -56,6 +67,7 @@ public function partitionByFields(string ...$fields): self
      * @param array<string, int|string>|string $fieldName Field name or array of field/order pairs
      * @param int|string                       $order     Field order (if one field is specified)
      * @psalm-param SortShape|string           $fieldName
+     * @psalm-param SortDirection|null         $order
      */
     public function sortBy($fieldName, $order = null): self
     {

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Group.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Group.php
@@ -9,6 +9,8 @@
 
 /**
  * Fluent interface for adding a $group stage to an aggregation pipeline.
+ *
+ * @psalm-type GroupStageExpression = array{'$group': array<string, mixed>}
  */
 class Group extends Operator
 {
@@ -22,6 +24,7 @@ public function __construct(Builder $builder)
         $this->expr = $builder->expr();
     }
 
+    /** @psalm-return GroupStageExpression */
     public function getExpression(): array
     {
         return [

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/IndexStats.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/IndexStats.php
@@ -9,9 +9,12 @@
 
 /**
  * Fluent interface for adding a $indexStats stage to an aggregation pipeline.
+ *
+ * @psalm-type IndexStatsStageExpression = array{'$indexStats': object}
  */
 class IndexStats extends Stage
 {
+    /** @psalm-return IndexStatsStageExpression */
     public function getExpression(): array
     {
         return [

diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Limit.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Limit.php
@@ -9,6 +9,8 @@
 
 /**
  * Fluent interface for adding a $limit stage to an aggregation pipeline.
+ *
+ * @psalm-type LimitStageExpression = array{'$limit': int}
  */
 class Limit extends Stage
 {
@@ -21,6 +23,7 @@ public function __construct(Builder $builder, int $limit)
         $this->limit = $limit;
     }
 
+    /** @psalm-return LimitStageExpression */
     public function getExpression(): array
     {
         return [