improving page/column usage

main.py

@@ -4,7 +4,5 @@ from icecream import ic
 
 ui.label('Hello, Nice GUI!')
 
-with ui.column():
-    ui.button('BUTTON 1', on_click=lambda _: ui.label('Nice!'))
-    ui.button('BUTTON 2', on_click=lambda _: ui.label('GUI!'))
-    
+with ui.column() as col:
+    col.button('BUTTON', on_click=lambda _: col.label('Nice!'))

nice_gui.py

@@ -4,25 +4,11 @@ from icecream import ic
 from contextlib import contextmanager
 
 
-class NiceGui():
-
-    def __init__(self):
-
-        self.context = [jp.WebPage(delete_flag=False)]
-
-        def build():
-            return self.context[0]
-
-        jp.justpy(build, start_server=False)
-
-    @property
-    def current(self):
-
-        return self.context[-1]
+class Group():
 
     def label(self, text):
 
-        p = jp.P(text=text, a=self.current, classes='w-48 text-xl p-1 m-2')
+        p = jp.P(text=text, a=self.jp, classes='flex text-xl p-1 m-2')
 
     def button(self, text, on_click=None):
 
@@ -32,21 +18,37 @@ class NiceGui():
             except:
                 traceback.print_exc()
 
-        d = jp.Div(text=text, a=self.current, classes='w-48 text-xl m-2 p-1 bg-blue-700 text-white text-center')
+        d = jp.Div(text=text, a=self.jp, classes='w-48 text-xl m-2 p-1 bg-blue-700 text-white text-center')
         d.on('click', click)
 
     @contextmanager
     def column(self):
-        d = jp.Div(a=self.current, classes='flex flex-wrap')
-        self.context.append(d)
-        yield
-        self.context.pop()
+        yield Column(self.jp)
+
+
+class Page(Group):
+
+    def __init__(self):
+
+        self.jp = jp.WebPage(delete_flag=False)
+
+        def build():
+            return self.jp
+
+        jp.justpy(build, start_server=False)
+
+
+class Column(Group):
+
+    def __init__(self, parent_jp) -> None:
+
+        self.jp = jp.Div(a=parent_jp, classes='flex flex-wrap')
 
 
-nice_gui = NiceGui()
+main = Page()
 ui = jp.app
 
 # bind methods to simplify API -- justpy creates an app which must be found by uvicorn via string "module:attribute"
-ui.label = nice_gui.label
-ui.button = nice_gui.button
-ui.column = nice_gui.column 
+ui.label = main.label
+ui.button = main.button
+ui.column = main.column