diff --git a/colab_logica.py b/colab_logica.py index afc760b..19dc02e 100755 --- a/colab_logica.py +++ b/colab_logica.py @@ -35,6 +35,7 @@ from IPython.core.magic import register_cell_magic from IPython.display import display +from IPython.display import HTML import os @@ -284,22 +285,14 @@ def __init__(self, bar, predicates): self.bar = bar self.predicates = predicates self.observed = set() + self.table_locations = {} def ObserveTable(self, predicate, table): - with self.bar.output_to(self.predicates.index(predicate)): - self.ObserveTableHere(predicate, table) + assert predicate in self.table_locations + self.table_locations[predicate].update(HTML(table.to_html())) - def ObserveTableHere(self, predicate, table): - if predicate in self.observed: - # Already observed. - return - self.observed |= {predicate} - print( - color.Format( - 'The following table is stored at {warning}%s{end} ' - 'variable.' % - predicate)) - display(table) + def RegisterTableLocation(self, predicate, table_location): + self.table_locations[predicate] = table_location def Logica(line, cell, run_query): @@ -344,6 +337,9 @@ def Logica(line, cell, run_query): executions = [] sub_bars = [] ip = IPython.get_ipython() + observer = None + if run_query: + observer = ExecutionObserver(bar, predicates) for idx, predicate in enumerate(predicates): with bar.output_to(logs_idx): try: @@ -373,6 +369,14 @@ def Logica(line, cell, run_query): else: print('Query is stored at %s variable.' % color.Warn(predicate + '_sql')) + with sub_bar.output_to(1): + print( + color.Format( + 'The following table is stored at {warning}%s{end} ' + 'variable.' % + predicate)) + table_location = display(HTML('table to be rendered'), display_id=True) + observer.RegisterTableLocation(predicate, table_location) with bar.output_to(logs_idx): if engine == 'sqlite': @@ -387,14 +391,11 @@ def Logica(line, cell, run_query): else: raise Exception('Logica only supports BigQuery, PostgreSQL and SQLite ' 'for now.') - observer = None - if run_query: - observer = ExecutionObserver(bar, predicates) try: result_map = concertina_lib.ExecuteLogicaProgram( executions, sql_runner=sql_runner, sql_engine=engine, display_mode=DISPLAY_MODE, - observer=None) # There are some errors caused by this. + observer=observer) except infer.TypeErrorCaughtException as e: e.ShowMessage() return @@ -405,7 +406,7 @@ def Logica(line, cell, run_query): with bar.output_to(idx): with sub_bars[idx].output_to(1): if run_query: - observer.ObserveTableHere(predicate, t) + observer.ObserveTable(predicate, t) else: print('The query was not run.') print(' ') # To activate the tabbar.