[Feature] Install necessary packages

README.md

@@ -9,9 +9,6 @@
 
 Run code server on Google Colab or Kaggle Notebooks
 
-## Getting Started
-
-
 ```shell
 # install colabcode
 $ pip install colabcode
@@ -26,16 +23,19 @@ $ ColabCode()
 # - port: the port you want to run code-server on, default 10000
 # - password: password to protect your code server from being accessed by someone else. Note that there is no password by default!
 # - mount_drive: True or False to mount your Google Drive
+# - packages: list of packages to be installed
 
-$ ColabCode(port=10000, password="abhishek", mount_drive=True)
+$ ColabCode(port=10000, password="abhishek", mount_drive=True, packages=['torchaudio', 'pytorch_lightning'])
 ```
 ## How to use it?
 Colab starter notebook: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abhishekkrthakur/colabcode/blob/master/colab_starter.ipynb)
 
 `ColabCode` comes pre-installed with some VS Code extensions.
 
+
 ##### See an example in youtube video     [![YouTube Video](https://img.shields.io/youtube/views/7kTbM3D02jU?style=social)](https://youtu.be/7kTbM3D02jU)
 
 ## License
 
 [MIT](LICENSE)
+

colabcode/code.py

@@ -1,7 +1,9 @@
 import os
 import subprocess
 from pyngrok import ngrok
-
+import sys
+import logging
+logging.basicConfig(filename='colabcode.log',level=logging.INFO)
 try:
     from google.colab import drive
 
@@ -14,15 +16,20 @@
 
 
 class ColabCode:
-    def __init__(self, port=10000, password=None, mount_drive=False):
+    def __init__(self, port=10000, password=None, mount_drive=False, packages=None):
+
         self.port = port
         self.password = password
         self._mount = mount_drive
         self._install_code()
         self._install_extensions()
         self._start_server()
+        self.packages = packages
+        if self.packages:
+            self._install_packages()
         self._run_code()
-
+
+
     def _install_code(self):
         subprocess.run(
             ["wget", "https://code-server.dev/install.sh"], stdout=subprocess.PIPE
@@ -58,3 +65,8 @@ def _run_code(self):
         ) as proc:
             for line in proc.stdout:
                 print(line, end="")
+
+    def _install_packages(self):
+        for package in self.packages:
+            logging.info(f'Installing {package}...')
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])