From 8ccf331988d937527b8ac72b929f2b7ba3cbbb2e Mon Sep 17 00:00:00 2001
From: Chris Campbell <135603296+c-campbell-mwam@users.noreply.github.com>
Date: Mon, 25 Mar 2024 09:50:49 +0000
Subject: [PATCH] fix: Allow single Java files to be added to context (#420)

---
 .../carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java b/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java
index 32140d0a..dd822278 100644
--- a/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java
+++ b/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java
@@ -40,6 +40,9 @@ public class PsiElementCheckboxTree extends FileCheckboxTree {
   }
 
   private static CheckedTreeNode createNode(PsiElement element) {
+    if (element instanceof com.intellij.psi.PsiClass) {
+      element = element.getContainingFile();
+    }
     if (!(element instanceof PsiDirectory || element instanceof PsiFile)) {
       return null;
     }