From 897f2578c6438842f47395641451a5ec53119664 Mon Sep 17 00:00:00 2001
From: Dmitrii Derepko <xepozz@list.ru>
Date: Wed, 14 Feb 2024 01:06:25 +0700
Subject: [PATCH] Add ability to use container as a factory

---
 src/Container.php                      | 30 +++++++++++++++++++++---
 src/Hook/AfterBuiltHook.php            | 21 +++++++++++++++++
 tests/Unit/Hook/AfterBuiltHookTest.php | 32 ++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 src/Hook/AfterBuiltHook.php
 create mode 100644 tests/Unit/Hook/AfterBuiltHookTest.php

diff --git a/src/Container.php b/src/Container.php
index 45065e0a..53cf3b48 100644
--- a/src/Container.php
+++ b/src/Container.php
@@ -35,7 +35,8 @@ final class Container implements ContainerInterface
 {
     private const META_TAGS = 'tags';
     private const META_RESET = 'reset';
-    private const ALLOWED_META = [self::META_TAGS, self::META_RESET];
+    private const META_AFTER_BUILT = 'afterBuilt';
+    private const ALLOWED_META = [self::META_TAGS, self::META_RESET, self::META_AFTER_BUILT];
 
     /**
      * @var DefinitionStorage Storage of object definitions.
@@ -67,6 +68,12 @@ final class Container implements ContainerInterface
      * @var Closure[]
      */
     private array $resetters = [];
+
+    /**
+     * @var Closure[]
+     */
+    public array $afterBuiltHooks = [];
+
     private bool $useResettersFromMeta = true;
 
     /**
@@ -134,10 +141,11 @@ public function has(string $id): bool
      */
     public function get(string $id)
     {
+        $result = null;
         if (!array_key_exists($id, $this->instances)) {
             try {
                 try {
-                    $this->instances[$id] = $this->build($id);
+                    $result = $this->instances[$id] = $this->build($id);
                 } catch (NotFoundExceptionInterface $e) {
                     if (!$this->delegates->has($id)) {
                         throw $e;
@@ -153,6 +161,7 @@ public function get(string $id)
                 throw new BuildingException($id, $e, $this->definitions->getBuildStack(), $e);
             }
         }
+        $this->callHookAfterBuilt($id);
 
         if ($id === StateResetter::class) {
             $delegatesResetter = null;
@@ -184,7 +193,7 @@ public function get(string $id)
         }
 
         /** @psalm-suppress MixedReturnStatement */
-        return $this->instances[$id];
+        return $result ?? $this->instances[$id];
     }
 
     /**
@@ -215,6 +224,9 @@ private function addDefinition(string $id, mixed $definition): void
         if (isset($meta[self::META_RESET])) {
             $this->setDefinitionResetter($id, $meta[self::META_RESET]);
         }
+        if (isset($meta[self::META_AFTER_BUILT])) {
+            $this->setAfterBuiltHook($id, $meta[self::META_AFTER_BUILT]);
+        }
 
         unset($this->instances[$id]);
         $this->addDefinitionToStorage($id, $definition);
@@ -621,4 +633,16 @@ private function buildProvider(mixed $provider): ServiceProviderInterface
 
         return $providerInstance;
     }
+
+    private function callHookAfterBuilt(string $id): void
+    {
+        if (isset($this->afterBuiltHooks[$id])) {
+            $this->afterBuiltHooks[$id]->call($this, $this, $id);
+        }
+    }
+
+    private function setAfterBuiltHook(string $id, Closure $callback): void
+    {
+        $this->afterBuiltHooks[$id] = $callback;
+    }
 }
diff --git a/src/Hook/AfterBuiltHook.php b/src/Hook/AfterBuiltHook.php
new file mode 100644
index 00000000..08e04cd8
--- /dev/null
+++ b/src/Hook/AfterBuiltHook.php
@@ -0,0 +1,21 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Di\Hook;
+
+use Closure;
+use Yiisoft\Di\Container;
+
+final class AfterBuiltHook
+{
+    public static function unsetInstance(): Closure
+    {
+        return function (Container $container, string $id) {
+            /**
+             * @var $this Container
+             */
+            unset($this->instances[$id]);
+        };
+    }
+}
diff --git a/tests/Unit/Hook/AfterBuiltHookTest.php b/tests/Unit/Hook/AfterBuiltHookTest.php
new file mode 100644
index 00000000..78e47321
--- /dev/null
+++ b/tests/Unit/Hook/AfterBuiltHookTest.php
@@ -0,0 +1,32 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Unit\Hook;
+
+use PHPUnit\Framework\TestCase;
+use Yiisoft\Di\Container;
+use Yiisoft\Di\ContainerConfig;
+use Yiisoft\Di\Hook\AfterBuiltHook;
+use Yiisoft\Di\Tests\Support\EngineInterface;
+use Yiisoft\Di\Tests\Support\EngineMarkOne;
+
+final class AfterBuiltHookTest extends TestCase
+{
+    public function testDifferentObjects()
+    {
+        $config = ContainerConfig::create()
+            ->withDefinitions([
+                EngineInterface::class => [
+                    'class' => EngineMarkOne::class,
+                    'afterBuilt' => AfterBuiltHook::unsetInstance()
+                ],
+            ]);
+        $container = new Container($config);
+
+        $this->assertNotSame(
+            $container->get(EngineInterface::class),
+            $container->get(EngineInterface::class),
+        );
+    }
+}