add from2d

src/fastpnm.nim

@@ -145,13 +145,18 @@ func binaryPosition(pnm: Pnm, x, y: Natural
 ): tuple[globalByteIndex, reverseBitIndex: int] =
     pnm.binaryPosition x+y*pnm.width
 
-func add*(pnm: var Pnm, v: bool) =
+func add*(pnm: var Pnm, b: bool) =
     let (gbi, rbi) = pnm.binaryPosition(pnm.filled)
     pnm.data.setlen gbi+1
     inc pnm.filled
-    if v:
+    if b:
         pnm.data[gbi].setBit(rbi)
 
+func add*(pnm: var Pnm, c: Color) =
+    pnm.data.add c.r
+    pnm.data.add c.g
+    pnm.data.add c.b
+
 const commonPnmExt* = ".pnm"
 func fileExt*(magic: PnmMagic): string =
     ## returns file extension according to the magic
@@ -214,24 +219,47 @@ func setColor*(pnm: var Pnm, x, y: int, color: Color) =
     pnm.data[i+2] = color.b
 
 func get2d*[T: bool or int or Color](pnm: Pnm): seq[seq[T]] =
+    ## generates 2D representaion of stored data
+    result.setLen pnm.height
+
     when T is bool:
         for y in 0..<pnm.height:
-            result.add @[]
             for x in 0..<pnm.width:
-                result[^1].add pnm.getBool(x, y)
+                result[y].add pnm.getBool(x, y)
 
     elif T is int:
         for y in 0..<pnm.height:
-            result.add @[]
             for x in 0..<pnm.width:
-                result[^1].add pnm.getGrayScale(x, y)
+                result[y].add pnm.getGrayScale(x, y)
 
     else:
         for y in 0..<pnm.height:
-            result.add @[]
             for x in 0..<pnm.width:
-                result[^1].add pnm.getColor(x, y)
+                result[y].add pnm.getColor(x, y)
+
+func from2d*[T: bool or int or Color](mat: seq[seq[T]], compress = true): Pnm =
+    let
+        h = mat.len
+        w = mat[0].len
+
+    result.width = w
+    result.height = h
+
+    result.magic =
+        when T is bool:
+            if compress: P4
+            else: P1
+        elif T is uint8:
+            if compress: P5
+            else: P2
+        else:
+            if compress: P6
+            else: P3
 
+    for y in 0..<h:
+        for x in 0..<w:
+            let t = mat[y][x]
+            result.add t
 
 # ----- main API
 

tests/test.nim

@@ -1,4 +1,4 @@
-import std/[unittest, os, strformat, sequtils]
+import std/[unittest, os, strformat, sequtils, strutils]
 import fastpnm
 
 
@@ -19,6 +19,8 @@ proc exportPnm(size: int, m: PnmMagic): string =
   fname
 
 func c(r, g, b: uint8): Color = (r, g, b)
+func `$`(c: Color): string =
+  c.r.toHex & c.g.toHex & c.b.toHex
 
 
 suite "raw & binary":
@@ -90,3 +92,19 @@ suite "special cases":
       img2 = parsePnm readfile "./examples/p3_spaces.ppm"
 
     checkSame img1, img2
+
+suite "other":
+  const
+    T = true
+    F = false
+
+  test "from2d":
+    let p = from2d @[
+      @[T, T, F, T, F, F, T, T, F, T],
+      @[F, F, T, F, T, T, F, F, T, T],
+      @[T, T, F, F, T, F, T, T, F, F]]
+
+    check cast[seq[uint8]](p.data) == @[
+      0b11010011'u8, 0b01000000'u8,
+      0b00101100'u8, 0b11000000'u8,
+      0b11001011'u8, 0b00000000'u8]