From e98bcbb1b1287553edd29fd7f3aeb68690fa7586 Mon Sep 17 00:00:00 2001
From: zhil <starsfm@gmail.com>
Date: Mon, 15 May 2017 10:04:55 +0300
Subject: [PATCH] added ability to limit doctrine collections - #10

---
 doc/usage.md                                  | 11 +++++
 .../FixtureGenerationContext.php              | 46 +++++++++++++++++++
 .../ObjectHandler/CollectionHandler.php       | 21 ++++++++-
 src/Trappar/AliceGenerator/ValueVisitor.php   |  9 ++++
 4 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/doc/usage.md b/doc/usage.md
index e12da74..7724aa6 100644
--- a/doc/usage.md
+++ b/doc/usage.md
@@ -32,15 +32,26 @@ Fixture Generation Contexts allow you to specify options which will affect a par
 ```php
 <?php
 use Trappar\AliceGenerator\FixtureGenerationContext;
+use Trappar\AliceGenerator\ObjectHandler\CollectionHandler;
+
+// this post will have 10 comments in fixtures
+CollectionHandler::limitCollection($post,"comments",10);
 
 // This will include only the Post
 $fixtureGenerator->generateYaml(
     $post,
     FixtureGenerationContext::create()
         ->setMaximumRecursion(2)
+        ->setMaximumCollectionChilds(5) // by default all collections would have 5 childs
+        ->setEntityCollectionLimit('AppBundle\Entity\Comment',3) // by default posts would have 3 comments
 );
 ```
 
+Collection limit check priorities 
+* CollectionHandler::limitCollection - you can limit some particular collection
+* ->setEntityCollectionLimit('AppBundle\Entity\Comment',3) - all Comment collections, which are not already limited with CollectionHandler::limitCollection
+* ->setMaximumCollectionChilds(5) - all not limited collections
+
 ### Limiting Recursion
 
 Since generating fixtures involves recursing through objects, you will often find yourself in a situation where generating fixtures yields quite a lot more information than you would like - this is just the nature of recursion! `FixtureGenerationContext` offers two solutions to this problem.
diff --git a/src/Trappar/AliceGenerator/FixtureGenerationContext.php b/src/Trappar/AliceGenerator/FixtureGenerationContext.php
index ef2e863..72a5b03 100644
--- a/src/Trappar/AliceGenerator/FixtureGenerationContext.php
+++ b/src/Trappar/AliceGenerator/FixtureGenerationContext.php
@@ -14,6 +14,17 @@ class FixtureGenerationContext
      * @var int
      */
     private $maximumRecursion = 5;
+
+    /**
+     * @var int
+     */
+    private $maximumCollectionChilds; // TODO: set to null to leave current behaviour. Maybe, lets set some default value, like 10?
+
+    /**
+     * @var array
+     */
+    private $entityCollectionLimits = [];
+
     /**
      * @var PersistedObjectConstraints
      */
@@ -143,4 +154,39 @@ public function setSortResults($sortResults)
 
         return $this;
     }
+
+    /**
+     * @return int
+     */
+    public function getMaximumCollectionChilds()
+    {
+        return $this->maximumCollectionChilds;
+    }
+
+    /**
+     * @param int $maximumCollectionChilds
+     * @return self
+     */
+    public function setMaximumCollectionChilds($maximumCollectionChilds)
+    {
+        $this->maximumCollectionChilds = $maximumCollectionChilds;
+        return $this;
+    }
+
+    public function getCollectionLimit($entityClassName)
+    {
+        if(isset($this->entityCollectionLimits[$entityClassName])) {
+            return $this->entityCollectionLimits[$entityClassName];
+        } elseif($this->maximumCollectionChilds) {
+            return $this->maximumCollectionChilds;
+        } else {
+            return false;
+        }
+    }
+
+    public function setEntityCollectionLimit($entityClassName,$limit)
+    {
+        $this->entityCollectionLimits[$entityClassName] = $limit;
+        return $this;
+    }
 }
\ No newline at end of file
diff --git a/src/Trappar/AliceGenerator/ObjectHandler/CollectionHandler.php b/src/Trappar/AliceGenerator/ObjectHandler/CollectionHandler.php
index 7580c3e..b5a2ace 100644
--- a/src/Trappar/AliceGenerator/ObjectHandler/CollectionHandler.php
+++ b/src/Trappar/AliceGenerator/ObjectHandler/CollectionHandler.php
@@ -15,10 +15,29 @@ public function handle(ValueContext $valueContext)
         if (!is_a($collection = $valueContext->getValue(), 'Doctrine\Common\Collections\Collection')) {
             return false;
         }
+        $fixturesGenerationContext = $valueContext->getValueVisitor()->getFixtureGenerationContext();
+        if(is_a($collection, 'Doctrine\ORM\PersistentCollection') && ($collectionLimit = $fixturesGenerationContext->getCollectionLimit($collection->getTypeClass()->getName()))) {
+            $refCollection = $valueContext->getMetadata()->reflection;
+            $criteria = \Doctrine\Common\Collections\Criteria::create()->setMaxResults($collectionLimit);
+            $refCollection->setAccessible(true);
+            $limitedValue = $collection->matching($criteria);
+            $refCollection->setValue($valueContext->getContextObject(), $limitedValue);
 
-        $valueContext->setValue($collection->toArray());
+            $valueContext->setValue($limitedValue);
+        } else {
+            $valueContext->setValue($collection->toArray());
+        }
         $valueContext->getValueVisitor()->visitArray($valueContext);
 
         return true;
     }
+
+    public static function limitCollection($entity, $collectionName, $limit)
+    {
+        $refMatchTeams = new \ReflectionProperty($entity,$collectionName);
+
+        $criteria = \Doctrine\Common\Collections\Criteria::create()->setMaxResults($limit);
+        $refMatchTeams->setAccessible(true);
+        $refMatchTeams->setValue($entity,$refMatchTeams->getValue($entity)->matching($criteria));
+    }
 }
\ No newline at end of file
diff --git a/src/Trappar/AliceGenerator/ValueVisitor.php b/src/Trappar/AliceGenerator/ValueVisitor.php
index 8e8f0db..6eddb78 100644
--- a/src/Trappar/AliceGenerator/ValueVisitor.php
+++ b/src/Trappar/AliceGenerator/ValueVisitor.php
@@ -243,4 +243,13 @@ private function handlePersistedObject($object, $reference)
             return true;
         }
     }
+
+    /**
+     * @return FixtureGenerationContext
+     */
+    public function getFixtureGenerationContext()
+    {
+        return $this->fixtureGenerationContext;
+    }
+
 }