From fa301cb48030eae3aa27b35925f3c80b49152683 Mon Sep 17 00:00:00 2001
From: Seongwon Im <imsw0529@gmail.com>
Date: Thu, 14 Sep 2023 15:09:06 +0900
Subject: [PATCH] [EdgeTpu] Add EdgeTpuConfigSetting.ts to separate config

This commit Add EdgeTpuConfigSetting.ts to separate ConfigObject logic of EdgeTpu Config

ONE-vscode-DCO-1.0-Signed-off-by: Seongwon Im <imsw0529@gmail.com>
---
 src/OneExplorer/ConfigObject.ts         |   7 ++
 src/OneExplorer/EdgeTpuConfigSetting.ts | 103 ++++++++++++++++++++++++
 src/OneExplorer/OneExplorer.ts          |   2 +-
 src/OneExplorer/OneStorage.ts           |   4 +-
 4 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 src/OneExplorer/EdgeTpuConfigSetting.ts

diff --git a/src/OneExplorer/ConfigObject.ts b/src/OneExplorer/ConfigObject.ts
index ce437ee78..3a0f797da 100644
--- a/src/OneExplorer/ConfigObject.ts
+++ b/src/OneExplorer/ConfigObject.ts
@@ -25,6 +25,7 @@ import { Logger } from "../Utils/Logger";
 
 import { Artifact } from "./ArtifactLocator";
 import { OneConfigSetting } from "./OneConfigSetting";
+import { EdgeTpuConfigSetting } from "./EdgeTpuConfigSetting";
 
 type Cfg = {
   "one-import-tflite": CfgOneImportTflite;
@@ -216,6 +217,9 @@ export class ConfigObj {
     const dir = path.dirname(filePath);
 
     let locatorRunner = OneConfigSetting.getBaseModelsLocatorRunner();
+    if (path.extname(filePath) === ".edgetpucfg") {
+      locatorRunner = EdgeTpuConfigSetting.getBaseModelsLocatorRunner();
+    }
 
     let artifacts: Artifact[] = locatorRunner.run(iniObj, dir);
 
@@ -254,6 +258,9 @@ export class ConfigObj {
     const dir = path.dirname(filePath);
 
     let locatorRunner = OneConfigSetting.getProductsLocatorRunner();
+    if (path.extname(filePath) === ".edgetpucfg") {
+      locatorRunner = EdgeTpuConfigSetting.getProductsLocatorRunner();
+    }
 
     /**
      * When you add a new product type, please append the ext type to
diff --git a/src/OneExplorer/EdgeTpuConfigSetting.ts b/src/OneExplorer/EdgeTpuConfigSetting.ts
new file mode 100644
index 000000000..d7c78a2d8
--- /dev/null
+++ b/src/OneExplorer/EdgeTpuConfigSetting.ts
@@ -0,0 +1,103 @@
+import * as vscode from "vscode";
+
+import { Locator, LocatorRunner } from "./ArtifactLocator";
+
+// export type OneCfg = {
+//   "one-import-tflite": CfgOneImportTflite;
+//   "one-import-onnx": CfgOneImportOnnx;
+//   "one-import-tf": CfgOneImportTf;
+// };
+// type CfgOneImportTflite = any;
+// type CfgOneImportOnnx = any;
+// type CfgOneImportTf = any;
+
+export class EdgeTpuConfigSetting {
+  static baseModelsLocatorRunner: LocatorRunner;
+  static productsLocatorRunner: LocatorRunner;
+  // TODO: make sections for updateBaseModelField method
+
+  static getBaseModelsLocatorRunner(): LocatorRunner {
+    if (!this.baseModelsLocatorRunner) {
+      this._initBaseModelsLocatorRunner();
+    }
+    return this.baseModelsLocatorRunner;
+  }
+
+  static getProductsLocatorRunner(): LocatorRunner {
+    if (!this.productsLocatorRunner) {
+      this._initProductsLocatorRunner();
+    }
+    return this.productsLocatorRunner;
+  }
+
+  private static _initBaseModelsLocatorRunner() {
+    let locatorRunner = new LocatorRunner();
+
+    locatorRunner.register({
+      artifactAttr: {
+        ext: ".tflite",
+        icon: new vscode.ThemeIcon("symbol-variable"),
+      },
+      locator: new Locator((value: string) => {
+        value += "";
+        const filterd = value
+          .split(" ")
+          .filter((val) => !val.endsWith("_edgetpu.tflite"));
+        value = filterd.join(" ");
+        return LocatorRunner.searchWithExt(".tflite", value);
+      }),
+    });
+
+    this.baseModelsLocatorRunner = locatorRunner;
+  }
+
+  private static _initProductsLocatorRunner() {
+    let locatorRunner = new LocatorRunner();
+
+    /**
+     * ABOUT ORDERING
+     *
+     * The registration order determines the order in the tree view
+     */
+
+    // NOTE
+    // Shows <model>_edgetpu.tflite
+    // <model>_edgetpu.tflite generated by <model>.tflite is product type
+    locatorRunner.register({
+      artifactAttr: {
+        ext: ".tflite",
+        icon: new vscode.ThemeIcon("symbol-variable"),
+      },
+      locator: new Locator((value: string) => {
+        value += "";
+        const filterd = value
+          .split(" ")
+          .filter((val) => val.endsWith("_edgetpu.tflite"));
+        value = filterd.join(" ");
+        return LocatorRunner.searchWithExt(".tflite", value);
+      }),
+    });
+
+    locatorRunner.register({
+      // 'default' view type is 'text editor' (vscode.openWith)
+      artifactAttr: {
+        ext: ".log",
+        openViewType: "default",
+        icon: vscode.ThemeIcon.File,
+        canHide: true,
+      },
+      locator: new Locator((value: string) => {
+        value += "";
+        const filterd = value
+          .split(" ")
+          .filter((val) => val.endsWith("_edgetpu.tflite"));
+        value = filterd.join(" ");
+        return LocatorRunner.searchWithExt(".tflite", value).map((filepath) =>
+          filepath.replace(".tflite", ".log")
+        );
+      }),
+    });
+
+    this.productsLocatorRunner = locatorRunner;
+  }
+}
diff --git a/src/OneExplorer/OneExplorer.ts b/src/OneExplorer/OneExplorer.ts
index 9b973c441..e6139aeb8 100644
--- a/src/OneExplorer/OneExplorer.ts
+++ b/src/OneExplorer/OneExplorer.ts
@@ -348,7 +348,7 @@ class BaseModelNode extends Node {
 class ConfigNode extends Node {
   readonly type = NodeType.config;
 
-  static readonly extList = [".cfg"];
+  static readonly extList = [".cfg", ".edgetpucfg"];
   // Open file with one.editor.cfg as default
   static defaultOpenViewType = "one.editor.cfg";
   // Display gear icon as default
diff --git a/src/OneExplorer/OneStorage.ts b/src/OneExplorer/OneStorage.ts
index a74826efc..5aa05466c 100644
--- a/src/OneExplorer/OneStorage.ts
+++ b/src/OneExplorer/OneStorage.ts
@@ -246,7 +246,9 @@ export class OneStorage {
     try {
       return roots
         .map((root) =>
-          readdirSyncRecursive(root).filter((val) => val.endsWith(".cfg"))
+          readdirSyncRecursive(root).filter(
+            (val) => val.endsWith(".cfg") || val.endsWith(".edgetpucfg")
+          )
         )
         .reduce((prev, cur) => [...prev, ...cur]);
     } catch {