From e681ef4d5268253925ec67b9eb419949ae79f9e7 Mon Sep 17 00:00:00 2001
From: Andrey Vokin <andrey.vokin@jetbrains.com>
Date: Fri, 31 Jan 2025 10:07:53 +0100
Subject: [PATCH] [python] cache PyImportElement#getImportReferenceExpression

GitOrigin-RevId: eccbb9acfc6d154c6e21b39128302c40e66107a3
---
 .../python/psi/impl/PyImportElementImpl.java       | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java
index 1d471547ac3a..47da215614ed 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java
@@ -3,6 +3,7 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.navigation.ItemPresentation;
+import com.intellij.openapi.util.Ref;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
@@ -26,6 +27,8 @@ import java.util.Objects;
  * The "import foo" or "import foo as bar" parts.
  */
 public class PyImportElementImpl extends PyBaseElementImpl<PyImportElementStub> implements PyImportElement {
+  private volatile @Nullable Ref<PyReferenceExpression> myImportReferenceExpression;
+
   public PyImportElementImpl(ASTNode astNode) {
     super(astNode);
   }
@@ -38,6 +41,17 @@ public class PyImportElementImpl extends PyBaseElementImpl<PyImportElementStub>
     super(stub, nodeType);
   }
 
+
+  @SuppressWarnings("DataFlowIssue")
+  @Override
+  public @Nullable PyReferenceExpression getImportReferenceExpression() {
+    PyReferenceExpression importReferenceExpression = myImportReferenceExpression == null ? null : myImportReferenceExpression.get();
+    if (myImportReferenceExpression == null || importReferenceExpression != null && !importReferenceExpression.isValid()) {
+      myImportReferenceExpression = new Ref<>(PyImportElement.super.getImportReferenceExpression());
+    }
+    return myImportReferenceExpression.get();
+  }
+
   @Override
   public QualifiedName getImportedQName() {
     final PyImportElementStub stub = getStub();