From 992d8d6856f2b3a323c92c78b4aeb3d60a117993 Mon Sep 17 00:00:00 2001
From: Cedric St-Jean <cedric.stjean@gmail.com>
Date: Sun, 20 Nov 2016 20:35:42 -0500
Subject: [PATCH] More flexible pyimport syntax. Fixes #66

---
 README.md        |  4 +++
 src/PyCall.jl    | 77 ++++++++++++++++++++++++++++++++++++++++++------
 test/runtests.jl |  3 ++
 3 files changed, 75 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 18a06f50..b2e4c9d2 100644
--- a/README.md
+++ b/README.md
@@ -117,6 +117,10 @@ compare it to the built-in Julia `sin`:
     @pyimport math
     math.sin(math.pi / 4) - sin(pi / 4)  # returns 0.0
 
+    # Alternative syntax
+    @pyimport math: (pow, radians)
+    pow(radians(180), 2)                 # returns pi²
+
 Type conversions are automatically performed for numeric, boolean,
 string, IO stream, date/period, and function types, along with tuples,
 arrays/lists, and dictionaries of these types. (Python functions can
diff --git a/src/PyCall.jl b/src/PyCall.jl
index ebc591b6..1a4be698 100644
--- a/src/PyCall.jl
+++ b/src/PyCall.jl
@@ -419,19 +419,78 @@ function pyimport_name(name, optional_varname)
     end
 end
 
-macro pyimport(name, optional_varname...)
-    mname = modulename(name)
-    Name = pyimport_name(name, optional_varname)
-    quote
-        if !isdefined($(Expr(:quote, Name)))
-            const $(esc(Name)) = pywrap(pyimport($mname))
-        elseif !isa($(esc(Name)), Module)
-            error("@pyimport: ", $(Expr(:quote, Name)), " already defined")
+# Helper for pyimport
+function pyimport_object_or_submodule(mod, member)
+    try
+        # Try to import member as a submodule of `mod`
+        return pywrap(pyimport(mod[:__name__] * "." * string(member)))
+    catch e
+        # Then it must be a variable/function
+        if isa(e, PyCall.PyError)
+            return mod[member]
+        else
+            rethrow()
         end
-        nothing
     end
 end
 
+"""
+`@pyimport` imports a Python module. Examples:
+
+```julia
+@pyimport scipy
+@pyimport scipy.stats: (betaprime, distributions)   # note the parentheses
+@pyimport scipy.stats as sstats
+
+println(scipy.average([1,2,3]))
+```
+"""
+macro pyimport(expr, optional_varname...)
+    if isa(expr, Expr)
+        # We could support `@pyimport module: a, b` but it's awkward because
+        # it parses as `((module:a), b)
+        @assert expr.head != :tuple "@pyimport requires parentheses, eg. @pyimport module_name: (a, b)"
+    end
+    if @capture(expr, mod_sym_:what_)
+        # Handle `@pyimport module_name: (a, b, c)`
+        @assert isempty(optional_varname) "Bad @pyimport syntax. See ?@pyimport"
+        if isa(what, Symbol)
+            members = [what]
+        else
+            @assert @capture(what, ((members__),)) "Bad @pyimport statement"
+        end
+        mod = gensym()
+        esc(quote
+            $mod = $PyCall.pyimport($(modulename(mod_sym)))
+            $([quote
+                  if !isdefined($(Expr(:quote, member)))
+                      const $member =
+                          $PyCall.pyimport_object_or_submodule($mod, $(Expr(:quote, member)))
+                  elseif !isa($member, Union{Module, PyObject})
+                      error("Cannot @pyimport ", $(Expr(:quote, member)),
+                            ". It is already defined in the current module.")
+                  end
+               end
+               for member in members]...)
+            nothing
+            end)
+    else
+        # Handle `@pyimport module_name` and `@pyimport module_name as bar`
+        mname = modulename(expr)
+        name = pyimport_name(expr, optional_varname)
+        quote
+            if !isdefined($(Expr(:quote, name)))
+                const $(esc(name)) = pywrap(pyimport($mname))
+            elseif !isa($(esc(name)), Module)
+                error("Cannot @pyimport ", $(Expr(:quote, name)),
+                      ". It is already defined in the current module.")
+            end
+            nothing
+        end
+    end
+end
+
+
 #########################################################################
 
 """
diff --git a/test/runtests.jl b/test/runtests.jl
index c6b62728..2d2e5928 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -98,6 +98,9 @@ array2py2arrayeq(x) = PyCall.py2array(Float64,PyCall.array2py(x)) == x
 @pyimport math
 @test_approx_eq math.sin(3) sin(3)
 
+@pyimport math: (pow, radians)
+@test_approx_eq pow(radians(180), 2) pi ^ 2
+
 @test collect(PyObject([1,"hello",5])) == [1,"hello",5]
 
 @test try @eval (@pyimport os.path) catch ex isa(ex, ArgumentError) end