From 2610d9e8b2f8a5356cb5bb85c400eb6945b478a6 Mon Sep 17 00:00:00 2001
From: Dimitri Lozeve <dimitri@lozeve.com>
Date: Wed, 4 Oct 2023 17:15:13 +0200
Subject: [PATCH] Allow to read and write more dtypes

---
 npy.bqn   | 29 +++++++++++++++++------------
 tests.bqn | 12 ++++++++++++
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/npy.bqn b/npy.bqn
index 7febe2a..477e6f7 100644
--- a/npy.bqn
+++ b/npy.bqn
@@ -1,40 +1,45 @@
 ⟨SaveNpy,LoadNpy⟩⇐
 
 lf←@+10
-magicString←(@+147)∾"NUMPY"  # "\x93NUMPY"
+magicString←(@+147)∾"NUMPY"
 
+# Type conversions
+dtypes‿formats←⟨"<f8","<i4","<u4"⟩‿⟨64‿'f',32‿'i',32‿'u'⟩
+DtypeToFormat←{(⊑dtypes⊐<𝕩)⊑formats}
+ArrayToBytes←{⟨DtypeToFormat 𝕨,8‿'c'⟩•bit._cast 𝕩}
+ArrayFromBytes←{⟨8‿'c',DtypeToFormat 𝕨⟩•bit._cast 𝕩}
+
+# Saving
 To_i16le←{@+256(|∾⌊∘÷˜)𝕩}
-# To_i16le←⟨16,8‿'c'⟩•bit._cast∘≍  # assumes the platform is little-endian
-FloatToBytes←⟨64,8‿'c'⟩•bit._cast
-
 FormatShape←{'('∾(1↓∾⥊','∾˘•Fmt¨≢𝕩)∾')'}
 
-BuildHeader←{𝕊𝕩:
+BuildHeader←{dtype 𝕊 data:
   version←1‿0
   # Hypothesis: BQN arrays are C-contiguous
-  headerData←"{'descr': '<f8', 'fortran_order': False, 'shape': "∾(FormatShape 𝕩)∾", }"
+  headerData←"{'descr':'"∾dtype∾"','fortran_order':False,'shape':"∾(FormatShape data)∾",}"
   padding←' '↑˜64-64|(≠magicString)+2+2+(≠headerData)+1
   ! 0=64|(≠magicString)+2+2+(≠padding)+(≠headerData)+1
   headerDataPadded←headerData∾padding∾lf
   magicString∾(@+version)∾(To_i16le ≠headerDataPadded)∾headerDataPadded
 }
-
-EncodeNpy←BuildHeader∾FloatToBytes∘⥊
-
+DetectDtype←{(∧´⥊⌊⊸=¨𝕩)⊑⟨"<f8",(∧´⥊0≤𝕩)⊑"<i4"‿"<u4"⟩}
+EncodeNpy←DetectDtype⊸(BuildHeader∾ArrayToBytes⟜⥊)
 SaveNpy←•file.Bytes⟜EncodeNpy
 
+# Loading
 ParseHeader←{𝕊𝕩:
   "Not a valid NPY file"! magicString≡(≠magicString)↑𝕩
   version←@-˜(0‿1+≠magicString)⊏𝕩
   headerlen←+´1‿256×@-˜(2‿3+≠magicString)⊏𝕩
   header←(¬∊⟜(' '∾lf))⊸/headerlen↑(4+≠magicString)↓𝕩
+  dtype←3↑(⊐⟜'<'↓⊢) (("descr"⍷header)⊐1)↓header
   shapestr←(⊐⟜')'↑⊢) (1⊸+∘⊐⟜'('↓⊢) (("shape"⍷header)⊐1)↓header
   shape←•BQN¨','((⊢-˜+`׬)∘=⊔⊢)shapestr
-  ⟨version,shape,headerlen+2+2+≠magicString⟩
+  ⟨version,dtype,shape,headerlen+2+2+≠magicString⟩
 }
 
 LoadNpy←{𝕊𝕩:
   bytes←•file.Bytes 𝕩
-  version‿shape‿datastart←ParseHeader bytes
-  shape⥊⟨8‿'c',64‿'f'⟩•bit._cast datastart↓bytes
+  version‿dtype‿shape‿datastart←ParseHeader bytes
+  shape⥊dtype ArrayFromBytes datastart↓bytes
 }
diff --git a/tests.bqn b/tests.bqn
index ae49f0f..d50906b 100644
--- a/tests.bqn
+++ b/tests.bqn
@@ -5,3 +5,15 @@ a←3‿4‿5•rand.Range 0
 b←LoadNpy"testbqn.npy"
 ! a≡○≢b
 ! a≡b
+
+floatArr←2‿3•rand.Range 0
+intArr←50-˜2‿3•rand.Range 100
+uintArr←2‿3•rand.Range 100
+
+"test_float.npy"SaveNpy floatArr
+"test_int.npy"SaveNpy intArr
+"test_uint.npy"SaveNpy uintArr
+
+floatArr≡LoadNpy"test_float.npy"
+intArr≡LoadNpy"test_int.npy"
+uintArr≡LoadNpy"test_uint.npy"