CODEWITHSUNDEEP
julia
Georgery
Georgery

Reputation: 8127

Function taking Vectors and Scalars

I have a function that takes a vector

function foo(x::Vector{Int64})
    x
end

How can I make it work also for scalars (i.e. turn them into one-element vectors)?

I know that I can do this:

foo(x::Int64) = foo([x])

but this stops being cool when there are more arguments, because you're writing multiple methods to achieve only one thing.

I think I something like foo(x::Union{Int64, Vector{Int64}}), but I don't know where or how it works or if it is the right thing to do.

Can anyone help?

Upvotes: 1

Views: 1025

Answers (2)

mcabbott
mcabbott

Reputation: 2590

You can make a helper function which either converts or does nothing. Then the main function can accept any combination:

_vec(x::Number) = [x]
_vec(x::AbstractVector) = x

function f(x, y, z)  # could specify ::Union{Number, AbstractVector}
  xv = _vec(x)
  yv = _vec(y)
  ...
end

The ... could do the actual work, or could call f(xv, yv, zv) where another method f(x::AbstractVector, y::AbstractVector, z::AbstractVector) does the work --- whichever seems cleaner.

Upvotes: 3

Oscar Smith
Oscar Smith

Reputation: 6398

The main time this comes up is if the version of your function for vectors does the same thing for all of it's elements. In this case, what you want to do is define f(x::Int), and use broadcasting f.([1,2,3]) for the vector case.

Upvotes: 1

Related Questions