From ff6287150370d7fd9be9d87c2adaf762ad71a16e Mon Sep 17 00:00:00 2001
From: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Date: Mon, 16 Oct 2023 09:11:26 +0200
Subject: [PATCH] fix requirement `lightning[app]`

---
 app.py                    |  4 ++--
 lai_jupyter/component.py  |  8 ++++----
 requirements.txt          |  6 +++---
 tests/test_jupyter_app.py | 10 +++++-----
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/app.py b/app.py
index aa1ef8e..f86845a 100644
--- a/app.py
+++ b/app.py
@@ -1,5 +1,5 @@
-import lightning as L
+from lightning.app import CloudCompute, LightningApp
 
 from lai_jupyter import JupyterLab
 
-app = L.LightningApp(JupyterLab(cloud_compute=L.CloudCompute("cpu-small")))
+app = LightningApp(JupyterLab(cloud_compute=CloudCompute("cpu-small")))
diff --git a/lai_jupyter/component.py b/lai_jupyter/component.py
index 8e8614e..506c1a6 100644
--- a/lai_jupyter/component.py
+++ b/lai_jupyter/component.py
@@ -6,7 +6,7 @@
 from time import sleep
 from typing import Literal, Optional
 
-import lightning as L
+from lightning.app import BuildConfig, CloudCompute, LightningWork
 
 R_INSTALL: list = """
 sudo apt-get update
@@ -28,7 +28,7 @@
 
 
 @dataclass
-class CustomBuildConfig(L.BuildConfig):
+class CustomBuildConfig(BuildConfig):
     kernel: str = "python"
 
     def build_commands(self) -> list:
@@ -36,11 +36,11 @@ def build_commands(self) -> list:
         return commands
 
 
-class JupyterLab(L.LightningWork):
+class JupyterLab(LightningWork):
     def __init__(
         self,
         kernel: Literal["python", "r", "julia"] = "python",
-        cloud_compute: Optional[L.CloudCompute] = None,
+        cloud_compute: Optional[CloudCompute] = None,
         parallel: bool = True,
     ) -> None:
         super().__init__(
diff --git a/requirements.txt b/requirements.txt
index c5609a2..359b3e9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-typing-extensions>=4.0.0
-jupyterlab==4.0.6
-lightning>=1.9.0, <3.0.0
+typing-extensions >=4.0.0
+jupyterlab ==4.0.6
+lightning[app] >=2.0.0, <2.2.0
 notebook
diff --git a/tests/test_jupyter_app.py b/tests/test_jupyter_app.py
index 0ed0c95..db31b77 100644
--- a/tests/test_jupyter_app.py
+++ b/tests/test_jupyter_app.py
@@ -1,13 +1,13 @@
 import os
 
-import lightning as L
 import requests
+from lightning.app import CloudCompute, LightningApp, LightningFlow
 from lightning.app.runners import MultiProcessRuntime
 
 from lai_jupyter import JupyterLab
 
 
-class TestJupyterServer(L.LightningFlow):
+class TestJupyterServer(LightningFlow):
     def __init__(self):
         super().__init__()
         self.success = False
@@ -18,11 +18,11 @@ def run(self, jupyter_url: str):
             self.success = True
 
 
-class RootFlow(L.LightningFlow):
+class RootFlow(LightningFlow):
     def __init__(self):
         super().__init__()
         self.jupyter_work = JupyterLab(
-            kernel="python", cloud_compute=L.CloudCompute(os.getenv("LIGHTNING_JUPYTER_LAB_COMPUTE", "cpu-small"))
+            kernel="python", cloud_compute=CloudCompute(os.getenv("LIGHTNING_JUPYTER_LAB_COMPUTE", "cpu-small"))
         )
         self.test_jupyter_flow = TestJupyterServer()
 
@@ -40,5 +40,5 @@ def configure_layout(self):
 
 
 def test_file_server():
-    app = L.LightningApp(RootFlow())
+    app = LightningApp(RootFlow())
     MultiProcessRuntime(app).dispatch()