From 7ef9fd4c8a50e95a5ae5897efdacac0c196479f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89tienne=20T=C3=A9treault-Pinard?= <etienne@plot.ly>
Date: Thu, 11 Jul 2019 15:52:41 -0400
Subject: [PATCH] fix clean2dArray for categorial x/y with extra items

---
 src/traces/heatmap/clean_2d_array.js |  6 +--
 test/jasmine/tests/heatmap_test.js   | 78 +++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/src/traces/heatmap/clean_2d_array.js b/src/traces/heatmap/clean_2d_array.js
index 13bc6432eeb..a5e0fd964e2 100644
--- a/src/traces/heatmap/clean_2d_array.js
+++ b/src/traces/heatmap/clean_2d_array.js
@@ -25,11 +25,11 @@ module.exports = function clean2dArray(zOld, trace, xa, ya) {
         for(i = 0; i < zOld.length; i++) rowlen = Math.max(rowlen, zOld[i].length);
         if(rowlen === 0) return false;
         getCollen = function(zOld) { return zOld.length; };
-        old2new = function(zOld, i, j) { return zOld[j][i]; };
+        old2new = function(zOld, i, j) { return (zOld[j] || [])[i]; };
     } else {
         rowlen = zOld.length;
         getCollen = function(zOld, i) { return zOld[i].length; };
-        old2new = function(zOld, i, j) { return zOld[i][j]; };
+        old2new = function(zOld, i, j) { return (zOld[i] || [])[j]; };
     }
 
     var padOld2new = function(zOld, i, j) {
@@ -58,9 +58,9 @@ module.exports = function clean2dArray(zOld, trace, xa, ya) {
     var xMap = axisMapping(xa);
     var yMap = axisMapping(ya);
 
+    if(ya && ya.type === 'category') rowlen = ya._categories.length;
     var zNew = new Array(rowlen);
 
-    if(ya && ya.type === 'category') rowlen = ya._categories.length;
     for(i = 0; i < rowlen; i++) {
         if(xa && xa.type === 'category') {
             collen = xa._categories.length;
diff --git a/test/jasmine/tests/heatmap_test.js b/test/jasmine/tests/heatmap_test.js
index aa85f05d7bd..875a92dc7f7 100644
--- a/test/jasmine/tests/heatmap_test.js
+++ b/test/jasmine/tests/heatmap_test.js
@@ -12,7 +12,6 @@ var destroyGraphDiv = require('../assets/destroy_graph_div');
 var supplyAllDefaults = require('../assets/supply_defaults');
 var failTest = require('../assets/fail_test');
 
-
 describe('heatmap supplyDefaults', function() {
     'use strict';
 
@@ -542,6 +541,83 @@ describe('heatmap calc', function() {
             expect(out.z[0][0]).toEqual(0);
         });
     });
+
+    describe('should clean z array linked to category x/y coordinates', function() {
+        var z = [
+            [1, 20, 30, 50, 1],
+            [20, 1, 60, 80, 30],
+            [30, 60, 1, -10, 20]
+        ];
+
+        it('- base case', function() {
+            var out = _calc({
+                z: z,
+                x: ['a', 'b', 'c', 'd', 'f'],
+                y: ['A', 'B', 'C']
+            });
+            expect(out.z).toBeCloseTo2DArray([
+                [1, 20, 30, 50, 1],
+                [20, 1, 60, 80, 30],
+                [30, 60, 1, -10, 20]
+            ]);
+        });
+
+        it('- with extra x items', function() {
+            var out = _calc({
+                z: z,
+                x: ['a', 'b', 'c', 'd', 'f', ''],
+                y: ['A', 'B', 'C']
+            });
+            expect(out.z).toBeCloseTo2DArray([
+                [1, 20, 30, 50, 1, undefined],
+                [20, 1, 60, 80, 30, undefined],
+                [30, 60, 1, -10, 20, undefined]
+            ]);
+        });
+
+        it('- with extra y items', function() {
+            var out = _calc({
+                z: z,
+                x: ['a', 'b', 'c', 'd', 'f'],
+                y: ['A', 'B', 'C', '']
+            });
+            expect(out.z).toBeCloseTo2DArray([
+                [1, 20, 30, 50, 1],
+                [20, 1, 60, 80, 30],
+                [30, 60, 1, -10, 20],
+                new Array(5)
+            ]);
+        });
+
+        it('- with extra x and y items', function() {
+            var out = _calc({
+                z: z,
+                x: ['a', 'b', 'c', 'd', 'f', ''],
+                y: ['A', 'B', 'C', '']
+            });
+            expect(out.z).toBeCloseTo2DArray([
+                [1, 20, 30, 50, 1, undefined],
+                [20, 1, 60, 80, 30, undefined],
+                [30, 60, 1, -10, 20, undefined],
+                new Array(6)
+            ]);
+        });
+
+        it('- transposed, with extra x and y items', function() {
+            var out = _calc({
+                transpose: true,
+                z: z,
+                x: ['a', 'b', 'c', 'd', 'f', ''],
+                y: ['A', 'B', 'C', '']
+            });
+            expect(out.z).toBeCloseTo2DArray([
+                [1, 20, 30, undefined, undefined, undefined],
+                [20, 1, 60, undefined, undefined, undefined],
+                [30, 60, 1, undefined, undefined, undefined],
+                [50, 80, -10, undefined, undefined, undefined]
+            ]);
+        });
+    });
 });
 
 describe('heatmap plot', function() {