#!/usr/bin/env python
# coding: utf-8

# # Making IPython magics from click functions
# 
# Two fun facts:
# 
# - click makes nice command-line APIs in Python
# - IPython uses magics to expose Python functions via a shell-like syntax
# 
# First, let's create the hello world click command from the click docs:

# In[1]:


import click

@click.command()
@click.option('--count', default=1, help='Number of greetings.')
@click.option('--name', prompt='Your name',
              help='The person to greet.')
def hello(count, name):
    """Simple program that greets NAME for a total of COUNT times."""
    for x in range(count):
        click.echo('Hello %s!' % click.style(name, fg='green'))


# which we can call:

# In[2]:


hello(['--count', '2'])


# Two things:
# 
# 1. it raises `SystemExit`, and
# 2. `min` is not green
# 
# The latter is because click checks if sys.stdin is a tty, which it isn't in the notebook,
# and only enables colors if it is.
# 
# The check is in `click.utils.should_strip_ansi`:

# In[3]:


s = click.style('Hello World!', fg='green')
s


# In[4]:


get_ipython().run_line_magic('pinfo2', 'click.utils.should_strip_ansi')


# In[5]:


click.echo(s)


# We can force click not to check this and then echo should produce colored output:

# In[6]:


try:
    from unittest import mock
except ImportError: # py2
    import mock

with mock.patch(
    'click.utils.should_strip_ansi',
    lambda *args, **kwargs: False
):
    click.echo(s)


# Next, we can write the transform to turn a click command into an IPython magic.
# An IPython magic is a Python function called where the rest of the line
# is passed as a simple string.
# The click function wants a list of arguments, `sys.argv`-style.
# 
# So we need to do the following things:
# 
# 1. split the line into arguments with [`shlex.split`](https://docs.python.org/3/library/shlex.html#shlex.split)
# 2. tell the command that it's name is `%magicname`
# 3. tell `click.echo` that it should always color output
# 4. catch `SystemExit` errors and turning them into IPython magic `UsageError`s
# 5. finally, register the magic via IPython's [`register_magic_function`](https://ipython.readthedocs.io/en/stable/api/generated/IPython.core.interactiveshell.html#IPython.core.interactiveshell.InteractiveShell.register_magic_function)

# In[7]:


import shlex
try:
    from unittest import mock
except ImportError: # py2
    import mock

from IPython.core.error import UsageError

def click_magic(click_command, name=None):
    if name is None:
        name = click_command.name

    def magic_func(line):
        args = shlex.split(line)
        # bypass click's check for whether colors should be enabled
        with mock.patch(
            'click.utils.should_strip_ansi',
            lambda *args, **kwargs: False
        ):
            try:
                click_command(
                    shlex.split(line),
                    prog_name='%' + name,
                )
            except SystemExit as e:
                if e.code != 0:
                    raise UsageError("Command exited with status=%s" % e.code)
    
    get_ipython().register_magic_function(magic_func, magic_name=name)


# In[8]:


click_magic(hello)


# In[9]:


get_ipython().run_line_magic('hello', '--help')


# In[10]:


get_ipython().run_line_magic('hello', '--count 2 --name myname')


# In[11]:


get_ipython().run_line_magic('hello', '--unrecognized')