feat: support for multiple columns

4 files changed, 95 insertions, 29 deletions

diff --git a/src/helpers/arrayFunctions.ts b/src/helpers/arrayFunctions.ts
index 7e78e29..234e135 100644
--- a/src/helpers/arrayFunctions.ts
+++ b/src/helpers/arrayFunctions.ts
@@ -42,15 +42,15 @@ export function range(
     throw new Error('step must be greater than 0');
   }
 
-  // If remove is specified, remove the values from the range
-  if (remove) {
-    for (let i = start; i <= end; i += step) {
+  // Generate an array from start to end
+  for (let i = start; i <= end; i += step) {
+    if (remove) {
       if (!remove.includes(i)) {
         rangeArray.push(i);
       }
+    } else {
+      rangeArray.push(i);
     }
   }
-
-  // return the range
   return rangeArray;
 }
diff --git a/src/index.ts b/src/index.ts
index e97d401..4bf04aa 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,5 +1,5 @@
 import { data, generateHeader, setHeader, header } from './lib/frame';
-import { getSingleColumnDetails } from './lib/locate';
+import { getMultipleColumnDetails, getSingleColumnDetails } from './lib/locate';
 import { show, head, tail } from './lib/display';
 import { getSize, info } from './lib/info';
 import { flatten } from './lib/data';
@@ -31,12 +31,10 @@ class Izuku {
     if (typeof column === 'number' || typeof column === 'string') {
       const izSampler = getSingleColumnDetails(this, column);
       return new Izuku(izSampler.rowd, izSampler.rowh);
-    } else if (Array.isArray(column)) {
-      // check if each element is an integer
+    } else if (Array.isArray(column) || isArrayOfType(column, 'string')) {
       if (isArrayOfType(column, 'number')) {
-        // TODO
-      } else if (isArrayOfType(column, 'string')) {
-        // TODO
+        const izSampler = getMultipleColumnDetails(this, column);
+        return new Izuku(izSampler.rowd, izSampler.rowh);
       } else {
         throw new Error('Columns must be an array of integers or column names');
       }
diff --git a/src/lib/locate.ts b/src/lib/locate.ts
index b176e07..55ba374 100644
--- a/src/lib/locate.ts
+++ b/src/lib/locate.ts
@@ -2,28 +2,75 @@ interface Izuku {
   rowdata: unknown[][];
   columns: string[];
 }
+
+/**
+ * checkIfColumnExists - check if a column exists in the frame
+ * @param frame: the frame to be checked
+ * @param column: the column index or name to be checked
+ * @returns true if the column exists
+ */
+export function checkIfColumnExists(
+  frame: Izuku,
+  column: number | string
+): boolean {
+  if (typeof column === 'number') {
+    return column < frame.columns.length;
+  } else if (typeof column === 'string') {
+    return frame.columns.includes(column);
+  } else {
+    throw new Error('column must be a number or a string');
+  }
+}
+
 /**
  * column returns a single column of the frame, option is either the column name or the column index
  * @param column: the column to be returned
  * @returns a single column of the frame as array of arrays
  */
 export function getSingleColumnDetails(iz: Izuku, column: number | string) {
-  let extractedColumn: any[] = [];
-  let columnName = '';
-  if (typeof column === 'string') {
-    extractedColumn = iz.rowdata.map((row) => {
-      return row[iz.columns.indexOf(column)];
-    });
-    columnName = column;
+  if (checkIfColumnExists(iz, column)) {
+    let extractedColumn: any[] = [];
+    let columnName = '';
+    if (typeof column === 'string') {
+      extractedColumn = iz.rowdata.map((row) => {
+        return row[iz.columns.indexOf(column)];
+      });
+      columnName = column;
+    }
+    if (typeof column === 'number') {
+      extractedColumn = iz.rowdata.map((row) => {
+        return row[column];
+      });
+      columnName = iz.columns[column]
+        ? iz.columns[column]
+        : `column${column + 1}`;
+    }
+    extractedColumn = extractedColumn.map((item) => [item]);
+    return { rowd: extractedColumn, rowh: [columnName] };
+  } else {
+    throw new Error(`Column ${column} does not exist`);
   }
-  if (typeof column === 'number') {
-    extractedColumn = iz.rowdata.map((row) => {
-      return row[column];
-    });
-    columnName = iz.columns[column]
-      ? iz.columns[column]
-      : `column${column + 1}`;
-  }
-  extractedColumn = extractedColumn.map((item) => [item]);
-  return { rowd: extractedColumn, rowh: [columnName] };
+}
+
+/**
+ * getMultipleColumnDetails - returns multiple columns of the frame
+ * @param columns: an array of column names or indexes
+ * @returns an array of arrays containing the columns
+ */
+export function getMultipleColumnDetails(
+  iz: Izuku,
+  columns: Array<number | string>
+) {
+  const extractedColumns: any[][] = [];
+  const columnNames: string[] = [];
+  columns.forEach((column) => {
+    const columnDetails = getSingleColumnDetails(iz, column);
+    extractedColumns.push(columnDetails.rowd);
+    columnNames.push(columnDetails.rowh[0]);
+  });
+  // transpose the columns
+  const transposedColumns = extractedColumns[0].map((col, i) => {
+    return extractedColumns.map((row) => row[i]);
+  });
+  return { rowd: transposedColumns, rowh: columnNames };
 }
diff --git a/tests/columns.test.ts b/tests/columns.test.ts
index 32fb6f1..09cb0f6 100644
--- a/tests/columns.test.ts
+++ b/tests/columns.test.ts
@@ -18,7 +18,28 @@ describe('columns.ts', () => {
         'Tiger',
         'Killjoy'
       ];
-      expect(frame.column('Name').flatten()).to.deep.equal(dataToExpect);
+      const names = frame.column('Name');
+      expect(names.flatten()).to.deep.equal(dataToExpect);
+      expect(names.columns).to.deep.equal(['Name']);
+    });
+  });
+  describe('Get multiple columns', () => {
+    it('should return multiple columns', () => {
+      const dataToExpect = [
+        [['Arthur'], ['Male']],
+        [['Betty'], ['Female']],
+        [['Victor'], ['Male']],
+        [['Dodger'], ['Male']],
+        [['Rayan'], ['Male']],
+        [['Skitley'], ['Female']],
+        [['Victoria'], ['Female']],
+        [['Tiger'], ['Male']],
+        [['Killjoy'], ['Female']]
+      ];
+      const headersToExpect = ['Name', 'Gender'];
+      const namesAndGenders = frame.column([0, 2]);
+      expect(namesAndGenders.rowdata).to.deep.equal(dataToExpect);
+      expect(namesAndGenders.columns).to.deep.equal(headersToExpect);
     });
   });
 });