From 4c696c8b21ac95ac9f40fddec11d132b06277455 Mon Sep 17 00:00:00 2001
From: Yuri Kuznetsov <yurikuzn@users.noreply.github.com>
Date: Mon, 8 Nov 2021 16:05:53 +0200
Subject: [PATCH] type fixes

---
 .../Extensions/EntityManagerReturnType.php    | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/dev/PHPStan/Extensions/EntityManagerReturnType.php b/dev/PHPStan/Extensions/EntityManagerReturnType.php
index ecd2ee4629..43ffb1bdfc 100644
--- a/dev/PHPStan/Extensions/EntityManagerReturnType.php
+++ b/dev/PHPStan/Extensions/EntityManagerReturnType.php
@@ -39,6 +39,8 @@ use PHPStan\Type\ObjectType;
 use PHPStan\Type\UnionType;
 use PHPStan\Type\NullType;
 use PHPStan\Type\Generic\GenericObjectType;
+use PHPStan\Type\Generic\TemplateGenericObjectType;
+
 use PhpParser\Node\Scalar\String_;
 
 use RuntimeException;
@@ -49,12 +51,16 @@ use Espo\ORM\EntityManager;
 use Espo\Core\Utils\Util;
 
 use Espo\ORM\Repository\RDBRepository;
+use Espo\ORM\Repository\Repository;
 
 class EntityManagerReturnType implements DynamicMethodReturnTypeExtension
 {
     private $supportedMethodNameList = [
         'getEntity',
+        'getNewEntity',
+        'createEntity',
         'getRDBRepository',
+        'getRepository',
     ];
 
     private $entityNamespaceList = [
@@ -84,10 +90,18 @@ class EntityManagerReturnType implements DynamicMethodReturnTypeExtension
             return $this->getGetEntity($methodReflection, $methodCall, $scope);
         }
 
+        if ($methodName === 'getNewEntity' || $methodName === 'createEntity') {
+            return $this->getGetEntityNotNull($methodReflection, $methodCall, $scope);
+        }
+
         if ($methodName === 'getRDBRepository') {
             return $this->getGetRDBRepository($methodReflection, $methodCall, $scope);
         }
 
+        if ($methodName === 'getRepository') {
+            return $this->getGetRepository($methodReflection, $methodCall, $scope);
+        }
+
         throw new RuntimeException("Not supported method.");
     }
 
@@ -116,6 +130,25 @@ class EntityManagerReturnType implements DynamicMethodReturnTypeExtension
         ]);
     }
 
+    private function getGetEntityNotNull(
+        MethodReflection $methodReflection,
+        MethodCall $methodCall,
+        Scope $scope
+    ): Type {
+
+        $value = $methodCall->args[0]->value;
+
+        if (!$value instanceof String_) {
+            return new ObjectType(Entity::class);
+        }
+
+        $entityType = $value->value;
+
+        $className = $this->findEntityClassName($entityType) ?? Entity::class;
+
+        return new ObjectType($className);
+    }
+
     private function findEntityClassName(string $entityType): ?string
     {
         foreach ($this->entityNamespaceList as $namespace) {
@@ -151,4 +184,27 @@ class EntityManagerReturnType implements DynamicMethodReturnTypeExtension
 
         return new ObjectType(RDBRepository::class);
     }
+
+    private function getGetRepository(
+        MethodReflection $methodReflection,
+        MethodCall $methodCall,
+        Scope $scope
+    ): Type {
+
+        $value = $methodCall->args[0]->value;
+
+        if (!$value instanceof ObjectType) {
+            return new ObjectType(Repository::class);
+        }
+
+        $entityType = $value->value;
+
+        $entityClassName = $this->findEntityClassName($entityType);
+
+        if ($entityClassName) {
+            return new TemplateGenericObjectType(Repository::class, [new ObjectType($entityClassName)]);
+        }
+
+        return new ObjectType(Repository::class);
+    }
 }