querybook icon indicating copy to clipboard operation
querybook copied to clipboard

fix: Improved Python print support for MultiIndex, NumPy arrays, etc

Open baumandm opened this issue 3 months ago • 0 comments

This PR makes a few improvements to the patched Python print function:

DataFrames

Before

Screenshot 2025-09-23 at 2 59 36 PM

After

Screenshot 2025-09-23 at 3 01 12 PM

NumPy Arrays

Before

Screenshot 2025-09-23 at 2 59 25 PM

After

Screenshot 2025-09-23 at 3 00 56 PM

Lists / Dicts with DataFrames / Arrays

Before

Screenshot 2025-09-23 at 2 59 46 PM

After

Screenshot 2025-09-23 at 3 01 34 PM

Unfortunately, the react-json-view component doesn't make it super clean to plug in a custom renderer if a DataFrame is embedded in a List, for example. So I'm just rendering it as JSON, which I think is still an improvement over the previous plain text format.

I have a semi-hacky solution for this if you're interested, but considering lists of DataFrames seems less common, I left that out.

Full Sample Script

Here's the full script I used for testing:

# ---------- 1. Primitives ----------
print("=== String ===")
print("Hello World")

print("\n=== Integer ===")
print(42)

print("\n=== Float ===")
print(3.14)

# ---------- 2. Simple Containers ----------
print("\n=== List ===")
simple_list = [1, 2, 3, "hello"]
print(simple_list)

print("\n=== Dict ===")
simple_dict = {"name": "John", "age": 25, "city": "NYC"}
print(simple_dict)

# ---------- 3. f-Strings ----------
print("\n=== f-String (list) ===")
print(f"List: {simple_list}")

print("\n=== f-String (dict) ===")
print(f"Dict: {simple_dict}")

# ---------- 4. NumPy ----------
import numpy as np

print("\n=== NumPy 1-D ===")
np_1d = np.array([1, 2, 3, 4, 5])
print(np_1d)

print("\n=== NumPy 2-D ===")
np_2d = np.arange(12).reshape(3, 4)
print(np_2d)

# ---------- 5. Pandas – Basic ----------
import pandas as pd

print("\n=== DataFrame (simple) ===")
df_simple = pd.DataFrame({"name": ["Alice", "Bob"], "age": [25, 30]})
print(df_simple)

# ---------- 6. Pandas – MultiIndex ----------
print("\n=== DataFrame (groupby MultiIndex) ===")
sales = pd.DataFrame({
    "region": ["North", "North", "South", "South"],
    "product": ["A", "B", "A", "B"],
    "sales": [100, 150, 200, 120],
})
grouped = sales.groupby("region").agg({"sales": ["sum", "mean"]})
print(grouped)

# ---------- 7. Pandas – Pivot Table ----------
print("\n=== Pivot Table ===")
pivot = sales.pivot_table(values="sales", index="region", columns="product", aggfunc="sum")
print(pivot)

# ---------- 8. Complex MultiIndex ----------
print("\n=== Triple-level MultiIndex ===")
big = pd.DataFrame({
    "region": ["N", "N", "N", "S", "S", "S"] * 2,
    "product": ["A", "A", "B", "A", "A", "B"] * 2,
    "quarter": ["Q1", "Q2"] * 6,
    "sales": range(12),
    "profit": range(12, 24),
})
triple = big.groupby(["region", "product", "quarter"]).agg({"sales": ["sum", "mean"], "profit": "sum"})
print(triple)

# ---------- 9. Containers of DataFrames ----------
print("\n=== List of DataFrames ===")
df_list = [df_simple, grouped]
print(df_list)

print("\n=== Dict of DataFrames ===")
df_dict = {"customers": df_simple, "sales_summary": grouped}
print(df_dict)

# ---------- 10. Mixed Bag ----------
print("\n=== Mixed List ===")
mixed = [df_simple, np_1d, "text", 42]
print(mixed)

baumandm avatar Sep 23 '25 19:09 baumandm