Add DirEntry.{is_dir}

palaviv · palaviv · commit 548f3f34c48f · 2019-04-21T12:04:12.000+03:00
diff --git a/tests/snippets/stdlib_os.py b/tests/snippets/stdlib_os.py
@@ -57,6 +57,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 FILE_NAME = "test1"
 FILE_NAME2 = "test2"
+FOLDER = "dir1"
 CONTENT = b"testing"
 CONTENT2 = b"rustpython"
 CONTENT3 = b"BOYA"
@@ -75,15 +76,21 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 	assert os.read(fd, len(CONTENT3)) == CONTENT3
 	os.close(fd)
 
-
 	fname2 = tmpdir + os.sep + FILE_NAME2
 	with open(fname2, "wb"):
 		pass
-	files = set()
+	folder = tmpdir + os.sep + FOLDER
+	os.mkdir(folder)
+
+	names = set()
 	paths = set()
+	dirs = set()
 	for dir_entry in os.scandir(tmpdir):
-		files.add(dir_entry.name)
+		names.add(dir_entry.name)
 		paths.add(dir_entry.path)
+		if dir_entry.is_dir():
+			dirs.add(dir_entry.name)
 
-	assert files == set([FILE_NAME, FILE_NAME2])
-	assert paths == set([fname, fname2])
+	assert names == set([FILE_NAME, FILE_NAME2, FOLDER])
+	assert paths == set([fname, fname2, folder])
+	assert dirs == set([FOLDER])
diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
@@ -214,6 +214,14 @@ impl DirEntryRef {
     fn path(self, _vm: &VirtualMachine) -> String {
         self.entry.path().to_str().unwrap().to_string()
     }
+
+    fn is_dir(self, vm: &VirtualMachine) -> PyResult<bool> {
+        Ok(self
+            .entry
+            .file_type()
+            .map_err(|s| vm.new_os_error(s.to_string()))?
+            .is_dir())
+    }
 }
 
 #[derive(Debug)]
@@ -275,6 +283,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let dir_entry = py_class!(ctx, "DirEntry", ctx.object(), {
          "name" => ctx.new_property(DirEntryRef::name),
          "path" => ctx.new_property(DirEntryRef::path),
+         "is_dir" => ctx.new_rustfunc(DirEntryRef::is_dir),
     });
 
     py_module!(vm, "_os", {
