Add support for JSON RPC requests

crates/pet-core/src/lib.rs

@@ -11,6 +11,7 @@ pub mod manager;
 pub mod os_environment;
 pub mod python_environment;
 pub mod reporter;
+// pub mod telemetry;
 
 #[derive(Debug, Clone)]
 pub struct LocatorResult {

crates/pet-core/src/reporter.rs

@@ -4,10 +4,7 @@
 use crate::{manager::EnvManager, python_environment::PythonEnvironment};
 
 pub trait Reporter: Send + Sync {
-    // fn get_reported_managers() -> Arc<Mutex<HashSet<PathBuf>>>;
-    // fn get_reported_environments() -> Arc<Mutex<HashSet<PathBuf>>>;
-
     fn report_manager(&self, manager: &EnvManager);
     fn report_environment(&self, env: &PythonEnvironment);
-    fn report_completion(&self, duration: std::time::Duration);
+    // fn report_telemetry(&self, event: &TelemetryEvent);
 }

crates/pet-homebrew/src/lib.rs

@@ -28,7 +28,7 @@ impl Homebrew {
     }
 }
 
-fn resolve(env: &PythonEnv) -> Option<PythonEnvironment> {
+fn from(env: &PythonEnv) -> Option<PythonEnvironment> {
     // Note: Sometimes if Python 3.10 was installed by other means (e.g. from python.org or other)
     // & then you install Python 3.10 via Homebrew, then some files will get installed via homebrew,
     // However everything (symlinks, Python executable `sys.executable`, `sys.prefix`) eventually point back to the existing installation.
@@ -83,7 +83,7 @@ fn resolve(env: &PythonEnv) -> Option<PythonEnvironment> {
 
 impl Locator for Homebrew {
     fn from(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
-        resolve(env)
+        from(env)
     }
 
     fn find(&self, reporter: &dyn Reporter) {
@@ -112,7 +112,7 @@ impl Locator for Homebrew {
                         // However this is a very generic location, and we might end up with other python installs here.
                         // Hence call `resolve` to correctly identify homebrew python installs.
                         let env_to_resolve = PythonEnv::new(file.clone(), None, None);
-                        if let Some(env) = resolve(&env_to_resolve) {
+                        if let Some(env) = from(&env_to_resolve) {
                             reporter.report_environment(&env);
                         }
                     });

crates/pet-jsonrpc/src/core.rs

@@ -28,3 +28,32 @@ pub fn send_message<T: serde::Serialize>(method: &'static str, params: Option<T>
     );
     let _ = io::stdout().flush();
 }
+pub fn send_reply<T: serde::Serialize>(id: u32, payload: Option<T>) {
+    let payload = serde_json::json!({
+        "jsonrpc": "2.0",
+        "result": payload,
+        "id": id
+    });
+    let message = serde_json::to_string(&payload).unwrap();
+    print!(
+        "Content-Length: {}\r\nContent-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n{}",
+        message.len(),
+        message
+    );
+    let _ = io::stdout().flush();
+}
+
+pub fn send_error(id: Option<u32>, code: i32, message: String) {
+    let payload = serde_json::json!({
+        "jsonrpc": "2.0",
+        "error": { "code": code, "message": message },
+        "id": id
+    });
+    let message = serde_json::to_string(&payload).unwrap();
+    print!(
+        "Content-Length: {}\r\nContent-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n{}",
+        message.len(),
+        message
+    );
+    let _ = io::stdout().flush();
+}

crates/pet-jsonrpc/src/lib.rs

@@ -2,7 +2,16 @@
 // Licensed under the MIT License.
 
 mod core;
+pub mod server;
 
 pub fn send_message<T: serde::Serialize>(method: &'static str, params: Option<T>) {
     core::send_message(method, params)
 }
+
+pub fn send_reply<T: serde::Serialize>(id: u32, payload: Option<T>) {
+    core::send_reply(id, payload)
+}
+
+pub fn send_error(id: Option<u32>, code: i32, message: String) {
+    core::send_error(id, code, message)
+}

crates/pet-jsonrpc/src/server.rs

@@ -0,0 +1,159 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::core::send_error;
+use serde_json::{self, Value};
+use std::{
+    collections::HashMap,
+    io::{self, Read},
+    sync::Arc,
+};
+
+type RequestHandler<C> = Arc<dyn Fn(Arc<C>, u32, Value)>;
+type NotificationHandler<C> = Arc<dyn Fn(Arc<C>, Value)>;
+
+pub struct HandlersKeyedByMethodName<C> {
+    context: Arc<C>,
+    requests: HashMap<&'static str, RequestHandler<C>>,
+    notifications: HashMap<&'static str, NotificationHandler<C>>,
+}
+
+impl<C> HandlersKeyedByMethodName<C> {
+    pub fn new(context: Arc<C>) -> Self {
+        HandlersKeyedByMethodName {
+            context,
+            requests: HashMap::new(),
+            notifications: HashMap::new(),
+        }
+    }
+
+    pub fn add_request_handler<F>(&mut self, method: &'static str, handler: F)
+    where
+        F: Fn(Arc<C>, u32, Value) + Send + Sync + 'static,
+    {
+        self.requests.insert(
+            method,
+            Arc::new(move |context, id, params| {
+                handler(context, id, params);
+            }),
+        );
+    }
+
+    pub fn add_notification_handler<F>(&mut self, method: &'static str, handler: F)
+    where
+        F: Fn(Arc<C>, Value) + Send + Sync + 'static,
+    {
+        self.notifications.insert(
+            method,
+            Arc::new(move |context, params| {
+                handler(context, params);
+            }),
+        );
+    }
+
+    fn handle_request(&self, message: Value) {
+        match message["method"].as_str() {
+            Some(method) => {
+                if let Some(id) = message["id"].as_u64() {
+                    if let Some(handler) = self.requests.get(method) {
+                        handler(self.context.clone(), id as u32, message["params"].clone());
+                    } else {
+                        eprint!("Failed to find handler for method: {}", method);
+                        send_error(
+                            Some(id as u32),
+                            -1,
+                            format!("Failed to find handler for request {}", method),
+                        );
+                    }
+                } else {
+                    // No id, so this is a notification
+                    if let Some(handler) = self.notifications.get(method) {
+                        handler(self.context.clone(), message["params"].clone());
+                    } else {
+                        eprint!("Failed to find handler for method: {}", method);
+                        send_error(
+                            None,
+                            -2,
+                            format!("Failed to find handler for notification {}", method),
+                        );
+                    }
+                }
+            }
+            None => {
+                eprint!("Failed to get method from message: {}", message);
+                send_error(
+                    None,
+                    -3,
+                    format!(
+                        "Failed to extract method from JSONRPC payload {:?}",
+                        message
+                    ),
+                );
+            }
+        };
+    }
+}
+
+/// Starts the jsonrpc server that listens for requests on stdin.
+/// This function will block forever.
+pub fn start_server<C>(handlers: &HandlersKeyedByMethodName<C>) -> ! {
+    let mut stdin = io::stdin();
+    loop {
+        let mut input = String::new();
+        match stdin.read_line(&mut input) {
+            Ok(_) => {
+                let mut empty_line = String::new();
+                match get_content_length(&input) {
+                    Ok(content_length) => {
+                        let _ = stdin.read_line(&mut empty_line);
+                        let mut buffer = vec![0; content_length];
+
+                        match stdin.read_exact(&mut buffer) {
+                            Ok(_) => {
+                                let request =
+                                    String::from_utf8_lossy(&buffer[..content_length]).to_string();
+                                match serde_json::from_str(&request) {
+                                    Ok(request) => handlers.handle_request(request),
+                                    Err(err) => {
+                                        eprint!("Failed to parse LINE: {}, {:?}", request, err)
+                                    }
+                                }
+                                continue;
+                            }
+                            Err(err) => {
+                                eprint!(
+                                    "Failed to read exactly {} bytes, {:?}",
+                                    content_length, err
+                                )
+                            }
+                        }
+                    }
+                    Err(err) => eprint!("Failed to get content length from {}, {:?}", input, err),
+                };
+            }
+            Err(error) => println!("Error in reading a line from stdin: {error}"),
+        }
+    }
+}
+
+/// Parses the content length from the given line.
+fn get_content_length(line: &str) -> Result<usize, String> {
+    let line = line.trim();
+    if let Some(content_length) = line.find("Content-Length: ") {
+        let start = content_length + "Content-Length: ".len();
+        if let Ok(length) = line[start..].parse::<usize>() {
+            Ok(length)
+        } else {
+            Err(format!(
+                "Failed to parse content length from {} for {}",
+                &line[start..],
+                line
+            ))
+        }
+    } else {
+        Err(format!(
+            "String 'Content-Length' not found in input => {}",
+            line
+        ))
+    }
+}

crates/pet-python-utils/src/headers.rs

@@ -42,22 +42,25 @@ pub fn get_version(path: &Path) -> Option<String> {
         if let Ok(result) = fs::read_to_string(patchlevel_h) {
             contents = result;
         } else if fs::metadata(&headers_path).is_err() {
+            // TODO: Remove this check, unnecessary, as we try to read the dir below.
             // Such a path does not exist, get out.
             continue;
         } else {
             // Try the other path
             // Sometimes we have it in a sub directory such as `python3.10` or `pypy3.9`
             if let Ok(readdir) = fs::read_dir(&headers_path) {
-                for path in readdir.filter_map(Result::ok).map(|e| e.path()) {
-                    if let Ok(metadata) = fs::metadata(&path) {
-                        if metadata.is_dir() {
-                            let patchlevel_h = path.join("patchlevel.h");
-                            if let Ok(result) = fs::read_to_string(patchlevel_h) {
-                                contents = result;
-                                break;
-                            }
+                for path in readdir.filter_map(Result::ok) {
+                    if let Ok(t) = path.file_type() {
+                        if !t.is_dir() {
+                            continue;
                         }
                     }
+                    let path = path.path();
+                    let patchlevel_h = path.join("patchlevel.h");
+                    if let Ok(result) = fs::read_to_string(patchlevel_h) {
+                        contents = result;
+                        break;
+                    }
                 }
             }
         }

crates/pet-reporter/src/jsonrpc.rs

@@ -39,9 +39,6 @@ impl Reporter for JsonRpcReporter {
             }
         }
     }
-    fn report_completion(&self, duration: std::time::Duration) {
-        send_message("exit", duration.as_millis().into())
-    }
 }
 
 pub fn create_reporter() -> impl Reporter {

crates/pet-reporter/src/stdio.rs

@@ -40,9 +40,6 @@ impl Reporter for StdioReporter {
             }
         }
     }
-    fn report_completion(&self, duration: std::time::Duration) {
-        println!("Refresh completed in {}ms", duration.as_millis())
-    }
 }
 
 pub fn create_reporter() -> impl Reporter {

crates/pet-reporter/src/test.rs

@@ -50,9 +50,6 @@ impl Reporter for TestReporter {
             }
         }
     }
-    fn report_completion(&self, _duration: std::time::Duration) {
-        //
-    }
 }
 
 pub fn create_reporter() -> TestReporter {

crates/pet/Cargo.toml

@@ -28,10 +28,10 @@ pet-pipenv = { path = "../pet-pipenv" }
 pet-global-virtualenvs = { path = "../pet-global-virtualenvs" }
 log = "0.4.21"
 clap = { version = "4.5.4", features = ["derive"] }
-
-[dev_dependencies]
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.93"
+
+[dev_dependencies]
 regex = "1.10.4"
 lazy_static = "1.4.0"