fc7627756370-5.py

# ipython-utils/researchenv/magics.py

from IPython.core.magic import Magics, magics_class, cell_magic, needs_local_scope
from IPython.core.magic_arguments import magic_arguments, argument, parse_argstring
from IPython.core.interactiveshell import InteractiveShell
from IPython.core.display import display
from google.cloud import bigquery

@magics_class
class ResearchEnvMagics(Magics):
    def __init__(self, shell: InteractiveShell, bq: bigquery.Client):
            super(ResearchEnvMagics, self).__init__(shell)
            self._bq = bq
      
    @magic_arguments()
    @argument('--name', help='Store result in a variable, otherwise '
                             'just prints the dataframe')
    @cell_magic
    def bq(self, line, cell):
        args = parse_argstring(self.bq, line)
        df = self._bq.query(cell).result().to_dataframe()
        if args.name is not None:
            self.shell.push({args.name: df})
        else:
            display(df)