use ruff_notebook::{Notebook, NotebookError};
use rustc_hash::FxHashMap;
use std::panic::RefUnwindSafe;
use std::sync::{Arc, Mutex};
use crate::Db;
use crate::files::File;
use crate::system::{
CommandExecutor, DirectoryEntry, MemoryFileSystem, Metadata, Result, System, SystemPath,
SystemPathBuf, SystemVirtualPath, WhichError, WhichResult,
};
use super::WritableSystem;
use super::walk_directory::WalkDirectoryBuilder;
#[derive(Debug)]
pub struct TestSystem {
inner: Arc<dyn WritableSystem + RefUnwindSafe + Send + Sync>,
env_overrides: Arc<Mutex<FxHashMap<String, Option<String>>>>,
}
impl Clone for TestSystem {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
env_overrides: self.env_overrides.clone(),
}
}
}
impl TestSystem {
pub fn new(inner: impl WritableSystem + RefUnwindSafe + Send + Sync + 'static) -> Self {
Self {
inner: Arc::new(inner),
env_overrides: Arc::new(Mutex::new(FxHashMap::default())),
}
}
pub fn set_env_var(&self, name: impl Into<String>, value: impl Into<String>) {
self.env_overrides
.lock()
.unwrap()
.insert(name.into(), Some(value.into()));
}
pub fn remove_env_var(&self, name: impl Into<String>) {
self.env_overrides.lock().unwrap().insert(name.into(), None);
}
pub fn in_memory(&self) -> &InMemorySystem {
self.as_in_memory()
.expect("The test db is not using a memory file system")
}
fn as_in_memory(&self) -> Option<&InMemorySystem> {
self.system().as_any().downcast_ref::<InMemorySystem>()
}
pub fn memory_file_system(&self) -> &MemoryFileSystem {
self.in_memory().fs()
}
fn use_system<S>(&mut self, system: S)
where
S: WritableSystem + Send + Sync + RefUnwindSafe + 'static,
{
self.inner = Arc::new(system);
}
fn system(&self) -> &dyn WritableSystem {
&*self.inner
}
}
impl System for TestSystem {
fn path_metadata(&self, path: &SystemPath) -> Result<Metadata> {
self.system().path_metadata(path)
}
fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf> {
self.system().canonicalize_path(path)
}
fn is_same_file(&self, first: &SystemPath, second: &SystemPath) -> Result<bool> {
self.system().is_same_file(first, second)
}
fn read_to_string(&self, path: &SystemPath) -> Result<String> {
self.system().read_to_string(path)
}
fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError> {
self.system().read_to_notebook(path)
}
fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
self.system().read_virtual_path_to_string(path)
}
fn read_virtual_path_to_notebook(
&self,
path: &SystemVirtualPath,
) -> std::result::Result<Notebook, NotebookError> {
self.system().read_virtual_path_to_notebook(path)
}
fn current_directory(&self) -> &SystemPath {
self.system().current_directory()
}
fn user_config_directory(&self) -> Option<SystemPathBuf> {
self.system().user_config_directory()
}
fn cache_dir(&self) -> Option<SystemPathBuf> {
self.system().cache_dir()
}
fn which(&self, _name: &str) -> WhichResult {
Err(WhichError::CannotFindBinaryPath)
}
fn command_executor(&self) -> Option<&dyn CommandExecutor> {
self.system().command_executor()
}
fn read_directory<'a>(
&'a self,
path: &SystemPath,
) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>> {
self.system().read_directory(path)
}
fn walk_directory(&self, path: &SystemPath) -> WalkDirectoryBuilder {
self.system().walk_directory(path)
}
fn as_writable(&self) -> Option<&dyn WritableSystem> {
Some(self)
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
fn env_var(&self, name: &str) -> std::result::Result<String, std::env::VarError> {
if let Some(override_value) = self.env_overrides.lock().unwrap().get(name) {
return match override_value {
Some(value) => Ok(value.clone()),
None => Err(std::env::VarError::NotPresent),
};
}
self.system().env_var(name)
}
fn dyn_clone(&self) -> Box<dyn System> {
Box::new(self.clone())
}
}
impl Default for TestSystem {
fn default() -> Self {
Self::new(InMemorySystem::default())
}
}
impl WritableSystem for TestSystem {
fn create_new_file(&self, path: &SystemPath) -> Result<()> {
self.system().create_new_file(path)
}
fn write_file_bytes(&self, path: &SystemPath, content: &[u8]) -> Result<()> {
self.system().write_file_bytes(path, content)
}
fn create_directory_all(&self, path: &SystemPath) -> Result<()> {
self.system().create_directory_all(path)
}
fn dyn_clone(&self) -> Box<dyn WritableSystem> {
Box::new(self.clone())
}
}
pub trait DbWithWritableSystem: Db + Sized {
type System: WritableSystem;
fn writable_system(&self) -> &Self::System;
fn write_file(&mut self, path: impl AsRef<SystemPath>, content: impl AsRef<str>) -> Result<()> {
let path = path.as_ref();
match self.writable_system().write_file(path, content.as_ref()) {
Ok(()) => {
File::sync_path(self, path);
Ok(())
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
if let Some(parent) = path.parent() {
self.writable_system().create_directory_all(parent)?;
for ancestor in parent.ancestors() {
File::sync_path(self, ancestor);
}
self.writable_system().write_file(path, content.as_ref())?;
File::sync_path(self, path);
Ok(())
} else {
Err(error)
}
}
err => err,
}
}
fn write_dedented(&mut self, path: &str, content: &str) -> Result<()> {
self.write_file(path, ruff_python_trivia::textwrap::dedent(content))?;
Ok(())
}
fn write_files<P, C, I>(&mut self, files: I) -> Result<()>
where
I: IntoIterator<Item = (P, C)>,
P: AsRef<SystemPath>,
C: AsRef<str>,
{
for (path, content) in files {
self.write_file(path, content)?;
}
Ok(())
}
}
pub trait DbWithTestSystem: Db + Sized {
fn test_system(&self) -> &TestSystem;
fn test_system_mut(&mut self) -> &mut TestSystem;
fn write_virtual_file(
&mut self,
path: impl AsRef<SystemVirtualPath>,
content: impl AsRef<[u8]>,
) {
let path = path.as_ref();
self.test_system()
.memory_file_system()
.write_virtual_file(path, content);
}
fn use_system<S>(&mut self, os: S)
where
S: WritableSystem + Send + Sync + RefUnwindSafe + 'static,
{
self.test_system_mut().use_system(os);
}
fn memory_file_system(&self) -> &MemoryFileSystem {
self.test_system().memory_file_system()
}
}
impl<T> DbWithWritableSystem for T
where
T: DbWithTestSystem,
{
type System = TestSystem;
fn writable_system(&self) -> &Self::System {
self.test_system()
}
}
#[derive(Clone, Default, Debug)]
pub struct InMemorySystem {
user_config_directory: Arc<Mutex<Option<SystemPathBuf>>>,
memory_fs: MemoryFileSystem,
}
impl InMemorySystem {
pub fn from_memory_fs(memory_fs: MemoryFileSystem) -> Self {
Self {
user_config_directory: Mutex::new(None).into(),
memory_fs,
}
}
pub fn fs(&self) -> &MemoryFileSystem {
&self.memory_fs
}
pub fn set_user_configuration_directory(&self, directory: Option<SystemPathBuf>) {
let mut user_directory = self.user_config_directory.lock().unwrap();
*user_directory = directory;
}
}
impl System for InMemorySystem {
fn path_metadata(&self, path: &SystemPath) -> Result<Metadata> {
self.memory_fs.metadata(path)
}
fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf> {
self.memory_fs.canonicalize(path)
}
fn is_same_file(&self, first: &SystemPath, second: &SystemPath) -> Result<bool> {
Ok(self.canonicalize_path(first)? == self.canonicalize_path(second)?)
}
fn read_to_string(&self, path: &SystemPath) -> Result<String> {
self.memory_fs.read_to_string(path)
}
fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError> {
let content = self.read_to_string(path)?;
Notebook::from_source_code(&content)
}
fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
self.memory_fs.read_virtual_path_to_string(path)
}
fn read_virtual_path_to_notebook(
&self,
path: &SystemVirtualPath,
) -> std::result::Result<Notebook, NotebookError> {
let content = self.read_virtual_path_to_string(path)?;
Notebook::from_source_code(&content)
}
fn current_directory(&self) -> &SystemPath {
self.memory_fs.current_directory()
}
fn user_config_directory(&self) -> Option<SystemPathBuf> {
self.user_config_directory.lock().unwrap().clone()
}
fn cache_dir(&self) -> Option<SystemPathBuf> {
None
}
fn which(&self, _name: &str) -> WhichResult {
Err(WhichError::CannotFindBinaryPath)
}
fn read_directory<'a>(
&'a self,
path: &SystemPath,
) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>> {
Ok(Box::new(self.memory_fs.read_directory(path)?))
}
fn walk_directory(&self, path: &SystemPath) -> WalkDirectoryBuilder {
self.memory_fs.walk_directory(path)
}
fn as_writable(&self) -> Option<&dyn WritableSystem> {
Some(self)
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
fn dyn_clone(&self) -> Box<dyn System> {
Box::new(self.clone())
}
}
impl WritableSystem for InMemorySystem {
fn create_new_file(&self, path: &SystemPath) -> Result<()> {
self.memory_fs.create_new_file(path)
}
fn write_file_bytes(&self, path: &SystemPath, content: &[u8]) -> Result<()> {
self.memory_fs.write_file(path, content)
}
fn create_directory_all(&self, path: &SystemPath) -> Result<()> {
self.memory_fs.create_directory_all(path)
}
fn dyn_clone(&self) -> Box<dyn WritableSystem> {
Box::new(self.clone())
}
}