Rewrite load.ps1 to support dot-sourcing and improve reliability

1 files changed, 34 insertions, 29 deletions

diff --git a/mling/tmpl/load.ps1 b/mling/tmpl/load.ps1
index d666338..fb616f4 100644
--- a/mling/tmpl/load.ps1
+++ b/mling/tmpl/load.ps1
@@ -3,14 +3,20 @@
 # Save original directory, restore after execution
 $_load_original_dir = Get-Location
 
-# Load completion script mling.ps1 from the current directory
-$mlingScript = Join-Path -Path (Get-Location) -ChildPath ".comp/mling_comp.ps1"
-if (Test-Path $mlingScript) {
-    . $mlingScript
+# Resolve script directory (works with dot-source: . ./load.ps1)
+$scriptPath = Split-Path -Parent $MyInvocation.MyCommand.Path
+if (-not $scriptPath) {
+    Write-Error "load.ps1: failed to resolve script directory"
+    return
+}
+
+# Load completion script mling_comp.ps1 from the .comp subdirectory
+$compScript = Join-Path -Path $scriptPath -ChildPath ".comp" | Join-Path -ChildPath "mling_comp.ps1"
+if (Test-Path $compScript) {
+    . $compScript
 }
 
 # Change to script directory
-$scriptPath = Split-Path -Parent $MyInvocation.MyCommand.Path
 try {
     Set-Location $scriptPath -ErrorAction Stop
 } catch {
@@ -19,59 +25,58 @@ try {
 }
 
 # Add bin directories from all namespaces to PATH
-Get-ChildItem -Directory -Path "*/bin/" | ForEach-Object {
-    $env:PATH = "$($_.FullName);$env:PATH"
+$allDirs = Get-ChildItem -Directory
+Get-ChildItem -Directory | Where-Object { Test-Path (Join-Path -Path $_.FullName -ChildPath "bin") } | ForEach-Object {
+    $binPath = Join-Path -Path $_.FullName -ChildPath "bin"
+    $env:PATH = "$binPath;$env:PATH"
 }
 
 # Helper function: execute script with appropriate shell
 function _load_script {
     param([string]$script)
-    # Only handle .ps1 scripts
     if ($script -like "*.ps1") {
         & $script 2>$null
     }
 }
 
-# Iterate over all namespaces
-Get-ChildItem -Directory | ForEach-Object {
-    $_namespace = $_.Name
+# Iterate over all namespaces (top-level directories except .comp)
+$nsDirs = Get-ChildItem -Directory -Exclude ".comp"
+foreach ($_dir in $nsDirs) {
+    $ns = $_dir.Name
 
     # Skip if UNTRUSTED marker exists
-    if (Test-Path "$_namespace\UNTRUSTED") { return }
+    $untrustedMarker = Join-Path -Path $ns -ChildPath "UNTRUSTED"
+    if (Test-Path $untrustedMarker) { continue }
 
-    $_comp_dir = "$_namespace\comp"
-    if (-not (Test-Path $_comp_dir -PathType Container)) { return }
+    $compDir = Join-Path -Path $ns -ChildPath "comp"
+    if (-not (Test-Path $compDir -PathType Container)) { continue }
 
     # Find all loadable scripts under comp
-    $_scripts = Get-ChildItem -Path $_comp_dir -File -Include "*.ps1" -ErrorAction SilentlyContinue
-    if (-not $_scripts) { return }
+    $scripts = Get-ChildItem -Path $compDir -Filter "*.ps1" -File -ErrorAction SilentlyContinue
+    if (-not $scripts) { continue }
 
-    # Count scripts
-    $_count = ($_scripts | Measure-Object).Count
+    $count = ($scripts | Measure-Object).Count
 
     # If TRUSTED marker exists, load directly
-    if (Test-Path "$_namespace\TRUSTED") {
-        $_scripts | ForEach-Object {
-            _load_script $_.FullName
-        }
-        return
+    $trustedMarker = Join-Path -Path $ns -ChildPath "TRUSTED"
+    if (Test-Path $trustedMarker) {
+        foreach ($_script in $scripts) { _load_script $_script.FullName }
+        continue
     }
 
     # No marker, ask user
-    $answer = Read-Host "'$_namespace' has $_count completion script(s) to load, do you trust it? [Y/n] "
+    $answer = Read-Host "'$ns' has $count completion script(s) to load, do you trust it? [Y/n] "
     if ($answer -eq "" -or $answer -match "^(y|yes)$") {
         # Mark as TRUSTED
-        New-Item -ItemType File -Path "$_namespace\TRUSTED" -Force | Out-Null
+        New-Item -ItemType File -Path $trustedMarker -Force | Out-Null
 
         # Ask whether to load immediately
         $load_answer = Read-Host "Load it immediately? [Y/n] "
         if ($load_answer -eq "" -or $load_answer -match "^(y|yes)$") {
-            $_scripts | ForEach-Object {
-                _load_script $_.FullName
-            }
+            foreach ($_script in $scripts) { _load_script $_script.FullName }
         }
     } else {
-        New-Item -ItemType File -Path "$_namespace\UNTRUSTED" -Force | Out-Null
+        New-Item -ItemType File -Path $untrustedMarker -Force | Out-Null
     }
 }