add size(A,d) methods

Author: bjarthur

src/TensorStoreWrapper.jl

@@ -150,6 +150,7 @@ function Base.setindex!(w::TensorStoreWrapper, v, indices...; kwargs...)
 end
 
 Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape)
+Base.size(w::TensorStoreWrapper, d::Integer) = d <= ndims(w) ? size(w)[d] : 1
 Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank)
 
 const TS_TYPE_MAP = Dict(
@@ -228,6 +229,7 @@ end
 
 # IndexDomainWrapper methods
 Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape)
+Base.size(w::IndexDomainWrapper, d::Integer) = d <= ndims(w) ? size(w)[d] : 1
 Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank)
 function Base.axes(w::IndexDomainWrapper)
     rank = ndims(w)

test/runtests.jl

@@ -23,6 +23,9 @@ using PythonCall
         @test eltype(w) == Int32
         @test ndims(w) == 2
         @test size(w) == (10, 20)
+        @test size(w, 1) == 10
+        @test size(w, 2) == 20
+        @test size(w, 3) == 1
         @test axes(w) == (1:10, 1:20)
     end
 
@@ -44,6 +47,9 @@ using PythonCall
         # Labeled indexing
         sub_w = w[x=1:5, y=11:15]
         @test size(sub_w) == (5, 5)
+        @test size(domain, 1) == 10
+        @test size(domain, 2) == 20
+        @test size(domain, 3) == 1
         @test axes(sub_w) == (1:5, 11:15)
 
         # translate_by