From 06567f6bfbf86d98904a7d32c2aa0a1a902012c7 Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Wed, 9 Sep 2026 04:59:23 +0800
Subject: [PATCH] Support PrettyTables 3 in sample display

---
 Project.toml     |  2 +-
 src/sampling.jl  |  6 +++++-
 test/sampling.jl | 12 +++++++++++-
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index a578ad1..23921e1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -25,7 +25,7 @@ DocStringExtensions = "0.8.6, 0.9"
 LinearAlgebra = "1"
 OMEinsum = "0.9.1"
 Pkg = "1"
-PrettyTables = "2"
+PrettyTables = "2, 3"
 ProblemReductions = "0.3"
 StatsBase = "0.34"
 TropicalNumbers = "0.5.4, 0.6"
diff --git a/src/sampling.jl b/src/sampling.jl
index 6c14c3c..88f447e 100644
--- a/src/sampling.jl
+++ b/src/sampling.jl
@@ -17,7 +17,11 @@ Base.length(s::Samples) = size(s.samples, 2)
 Base.size(s::Samples) = (size(s.samples, 2),)
 function Base.show(io::IO, s::Samples)  # display with PrettyTables
     println(io, typeof(s))
-    PrettyTables.pretty_table(io, s.samples', header=s.labels)
+    if pkgversion(PrettyTables) < v"3"
+        PrettyTables.pretty_table(io, s.samples'; header=s.labels)
+    else
+        PrettyTables.pretty_table(io, s.samples'; column_labels=s.labels)
+    end
 end
 num_samples(samples::Samples) = size(samples.samples, 2)
 function idx4labels(totalset::AbstractVector{L}, labels::AbstractVector{L})::Vector{Int} where L
diff --git a/test/sampling.jl b/test/sampling.jl
index a586984..625278f 100644
--- a/test/sampling.jl
+++ b/test/sampling.jl
@@ -2,6 +2,16 @@ using TensorInference, Test, LinearAlgebra
 import StatsBase
 using OMEinsum, Random
 
+@testset "Samples display" begin
+    samples = TensorInference.Samples([0 1; 1 0], ["x", "y"])
+    output = sprint(show, samples)
+    @test occursin("Samples{String}", output)
+    @test occursin("x", output)
+    @test occursin("y", output)
+    @test occursin("0", output)
+    @test occursin("1", output)
+end
+
 @testset "sampling" begin
     model = TensorInference.read_model_from_string("""MARKOV
 8
@@ -96,4 +106,4 @@ end
     mps.tensors[setdiff(1:length(mps.tensors), mps.unity_tensors_idx)] .*= 100
     samples = sample(mps, 1; rescale = true)
     @test samples isa TensorInference.Samples
-end
\ No newline at end of file
+end