refactor(vscode): refactor notebook uri conversion. (#3738)

TabbyML · Jan 21, 2025 · 812a622 · 812a622
1 parent 3875aee
commit 812a622
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 79 deletions.
diff --git a/clients/vscode/src/chat/utils.ts b/clients/vscode/src/chat/utils.ts
@@ -1,13 +1,12 @@
 import path from "path";
-import { Position as VSCodePosition, Range as VSCodeRange, Uri, workspace, TextEditor } from "vscode";
+import { TextEditor, Position as VSCodePosition, Range as VSCodeRange, Uri, workspace } from "vscode";
 import type {
   Filepath,
   Position as ChatPanelPosition,
   LineRange,
   PositionRange,
   Location,
   ListFileItem,
-  FilepathInGitRepository,
 } from "tabby-chat-panel";
 import type { GitProvider } from "../git/GitProvider";
 import { getLogger } from "../logger";
@@ -32,15 +31,12 @@ export function isValidForSyncActiveEditorSelection(editor: TextEditor): boolean
 }
 
 export function localUriToChatPanelFilepath(uri: Uri, gitProvider: GitProvider): Filepath {
-  let uriFilePath = uri.toString(true);
-  if (uri.scheme === DocumentSchemes.vscodeNotebookCell) {
-    const notebook = parseNotebookCellUri(uri);
-    if (notebook) {
-      // add fragment `#cell={number}` to filepath
-      uriFilePath = uri.with({ scheme: notebook.notebook.scheme, fragment: `cell=${notebook.handle}` }).toString(true);
-    }
+  let localUri = uri;
+  if (localUri.scheme === DocumentSchemes.vscodeNotebookCell) {
+    localUri = convertFromNotebookCellUri(localUri);
   }
 
+  const uriFilePath = uri.toString(true);
   const workspaceFolder = workspace.getWorkspaceFolder(uri);
 
   let repo = gitProvider.getRepository(uri);
@@ -65,80 +61,34 @@ export function localUriToChatPanelFilepath(uri: Uri, gitProvider: GitProvider):
   };
 }
 
-function isJupyterNotebookFilepath(filepath: Filepath): boolean {
-  const _filepath = filepath.kind === "uri" ? filepath.uri : filepath.filepath;
-  const extname = path.extname(_filepath);
-  return extname.startsWith(".ipynb");
-}
-
 export function chatPanelFilepathToLocalUri(filepath: Filepath, gitProvider: GitProvider): Uri | null {
-  const isNotebook = isJupyterNotebookFilepath(filepath);
-
+  let result: Uri | null = null;
   if (filepath.kind === "uri") {
     try {
-      if (isNotebook) {
-        const handle = chatPanelFilePathToNotebookCellHandle(filepath.uri);
-        if (typeof handle === "number") {
-          return generateLocalNotebookCellUri(Uri.parse(filepath.uri), handle);
-        }
-      }
-
-      return Uri.parse(filepath.uri, true);
+      result = Uri.parse(filepath.uri, true);
     } catch (e) {
       // FIXME(@icycodes): this is a hack for uri is relative filepaths in workspaces
       const workspaceRoot = workspace.workspaceFolders?.[0];
       if (workspaceRoot) {
-        return Uri.joinPath(workspaceRoot.uri, filepath.uri);
+        result = Uri.joinPath(workspaceRoot.uri, filepath.uri);
       }
     }
   } else if (filepath.kind === "git") {
     const localGitRoot = gitProvider.findLocalRootUriByRemoteUrl(filepath.gitUrl);
     if (localGitRoot) {
-      // handling for Jupyter Notebook (.ipynb) files
-      if (isNotebook) {
-        return chatPanelFilepathToVscodeNotebookCellUri(localGitRoot, filepath);
-      }
-
-      return Uri.joinPath(localGitRoot, filepath.filepath);
+      result = Uri.joinPath(localGitRoot, filepath.filepath);
     }
   }
-  logger.warn(`Invalid filepath params.`, filepath);
-  return null;
-}
-
-function chatPanelFilepathToVscodeNotebookCellUri(root: Uri, filepath: FilepathInGitRepository): Uri | null {
-  if (filepath.kind !== "git") {
-    logger.warn(`Invalid filepath params.`, filepath);
-    return null;
-  }
 
-  const filePathUri = Uri.parse(filepath.filepath);
-  const notebookUri = Uri.joinPath(root, filePathUri.path);
-
-  const handle = chatPanelFilePathToNotebookCellHandle(filepath.filepath);
-  if (typeof handle === "undefined") {
+  if (result == null) {
     logger.warn(`Invalid filepath params.`, filepath);
     return null;
   }
-  return generateLocalNotebookCellUri(notebookUri, handle);
-}
-
-function chatPanelFilePathToNotebookCellHandle(filepath: string): number | undefined {
-  let handle: number | undefined;
-
-  const fileUri = Uri.parse(filepath);
-  const fragment = fileUri.fragment;
-  const searchParams = new URLSearchParams(fragment);
-  if (searchParams.has("cell")) {
-    const cellString = searchParams.get("cell")?.toString() || "";
-    handle = parseInt(cellString, 10);
-  }
 
-  if (typeof handle === "undefined" || isNaN(handle)) {
-    return undefined;
+  if (isJupyterNotebookFilepath(result)) {
+    result = convertToNotebookCellUri(result);
   }
-
-  return handle;
+  return result;
 }
 
 export function vscodePositionToChatPanelPosition(position: VSCodePosition): ChatPanelPosition {
@@ -192,10 +142,51 @@ export function chatPanelLocationToVSCodeRange(location: Location | undefined):
   return null;
 }
 
+export function localUriToListFileItem(uri: Uri, gitProvider: GitProvider): ListFileItem {
+  return {
+    filepath: localUriToChatPanelFilepath(uri, gitProvider),
+  };
+}
+
+export function escapeGlobPattern(query: string): string {
+  // escape special glob characters: * ? [ ] { } ( ) ! @
+  return query.replace(/[*?[\]{}()!@]/g, "\\$&");
+}
+
+// Notebook cell uri conversion
+
+function isJupyterNotebookFilepath(uri: Uri): boolean {
+  const extname = path.extname(uri.fsPath);
+  return extname.startsWith(".ipynb");
+}
+
+function convertToNotebookCellUri(uri: Uri): Uri {
+  let handle: number | undefined;
+
+  const searchParams = new URLSearchParams(uri.fragment);
+  const cellString = searchParams.get("cell");
+  if (cellString) {
+    handle = parseInt(cellString, 10);
+  }
+  handle = handle || 0;
+
+  searchParams.set("cell", handle.toString());
+  return generateNotebookCellUri(uri, handle);
+}
+
+function convertFromNotebookCellUri(uri: Uri): Uri {
+  const parsed = parseNotebookCellUri(uri);
+  if (!parsed) {
+    return uri;
+  }
+  return uri.with({ scheme: parsed.notebook.scheme, fragment: `cell=${parsed.handle}` });
+}
+
 const nb_lengths = ["W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f"];
 const nb_padRegexp = new RegExp(`^[${nb_lengths.join("")}]+`);
 const nb_radix = 7;
-export function parseNotebookCellUri(cell: Uri): { notebook: Uri; handle: number } | undefined {
+
+function parseNotebookCellUri(cell: Uri): { notebook: Uri; handle: number } | undefined {
   if (cell.scheme !== DocumentSchemes.vscodeNotebookCell) {
     return undefined;
   }
@@ -216,20 +207,9 @@ export function parseNotebookCellUri(cell: Uri): { notebook: Uri; handle: number
   };
 }
 
-export function generateLocalNotebookCellUri(notebook: Uri, handle: number): Uri {
+function generateNotebookCellUri(notebook: Uri, handle: number): Uri {
   const s = handle.toString(nb_radix);
   const p = s.length < nb_lengths.length ? nb_lengths[s.length - 1] : "z";
   const fragment = `${p}${s}s${Buffer.from(notebook.scheme).toString("base64")}`;
   return notebook.with({ scheme: DocumentSchemes.vscodeNotebookCell, fragment });
 }
-
-export function uriToListFileItem(uri: Uri, gitProvider: GitProvider): ListFileItem {
-  return {
-    filepath: localUriToChatPanelFilepath(uri, gitProvider),
-  };
-}
-
-export function escapeGlobPattern(query: string): string {
-  // escape special glob characters: * ? [ ] { } ( ) ! @
-  return query.replace(/[*?[\]{}()!@]/g, "\\$&");
-}
diff --git a/clients/vscode/src/chat/webview.ts b/clients/vscode/src/chat/webview.ts
@@ -47,7 +47,7 @@ import {
   vscodeRangeToChatPanelPositionRange,
   chatPanelLocationToVSCodeRange,
   isValidForSyncActiveEditorSelection,
-  uriToListFileItem,
+  localUriToListFileItem,
   escapeGlobPattern,
 } from "./utils";
 import mainHtml from "./html/main.html";
@@ -509,7 +509,7 @@ export class ChatWebview {
             .filter((tab) => tab.input && (tab.input as TabInputText).uri);
 
           this.logger.info(`No query provided, listing ${openTabs.length} opened editors.`);
-          return openTabs.map((tab) => uriToListFileItem((tab.input as TabInputText).uri, this.gitProvider));
+          return openTabs.map((tab) => localUriToListFileItem((tab.input as TabInputText).uri, this.gitProvider));
         }
 
         try {
@@ -529,7 +529,7 @@ export class ChatWebview {
 
           const files = await workspace.findFiles(globPattern, undefined, maxResults);
           this.logger.info(`Found ${files.length} files.`);
-          return files.map((uri) => uriToListFileItem(uri, this.gitProvider));
+          return files.map((uri) => localUriToListFileItem(uri, this.gitProvider));
         } catch (error) {
           this.logger.warn("Failed to find files:", error);
           window.showErrorMessage("Failed to find files.");