From baafb2e887a9ff9a6284d70103db945ed53aed9b Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Tue, 20 Aug 2019 09:53:37 -0700
Subject: [PATCH] Updated new visualizations to match syntax of #1878

---
 backend/src/apiserver/visualization/table.py | 17 +++++++++--------
 backend/src/apiserver/visualization/tfma.py  |  6 +++---
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/backend/src/apiserver/visualization/table.py b/backend/src/apiserver/visualization/table.py
index f0cb828601b..b95ad22a5de 100644
--- a/backend/src/apiserver/visualization/table.py
+++ b/backend/src/apiserver/visualization/table.py
@@ -22,24 +22,25 @@
 import pandas as pd
 from tensorflow.python.lib.io import file_io
 
-# Remove maxByte limit
+# Remove maxByte limit to prevent issues where entire table cannot be rendered
+# due to size of data.
 opts.maxBytes = 0
 
 dfs = []
 files = file_io.get_matching_files(source)
 
 # Read data from file and write it to a DataFrame object.
-if "headers" is in variables:
-    # If headers are provided, do not set headers for DataFrames
-    for f in files:
-        dfs.append(pd.read_csv(f, header=None))
-else:
+if variables.get("headers", False) is False:
     # If no headers are provided, use the first row as headers
     for f in files:
         dfs.append(pd.read_csv(f))
+else:
+    # If headers are provided, do not set headers for DataFrames
+    for f in files:
+        dfs.append(pd.read_csv(f, header=None))
 
 # Display DataFrame as output.
 df = pd.concat(dfs)
-if "headers" is in variables:
-    df.columns = variables["headers"]
+if variables.get("headers", False) != False:
+    df.columns = variables.get("headers")
 show(df)
diff --git a/backend/src/apiserver/visualization/tfma.py b/backend/src/apiserver/visualization/tfma.py
index 647a1a700e3..e5ca988dc4f 100644
--- a/backend/src/apiserver/visualization/tfma.py
+++ b/backend/src/apiserver/visualization/tfma.py
@@ -20,9 +20,9 @@
 #
 # source
 
-if "slicing_column" in variables {
-    tfma.view.render_slicing_metrics(source, slicing_column=variables["slicing_column"])
-} else {
+if variables.get("slicing_column", False) is False {
     tfma.view.render_slicing_metrics(source)
+} else {
+    tfma.view.render_slicing_metrics(source, slicing_column=variables.get("slicing_column"))
 }