From 548f3f34c48ff55e63223fdb1a99f44fbae627fb Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Sun, 21 Apr 2019 12:04:12 +0300
Subject: [PATCH] Add DirEntry.{is_dir}

---
 tests/snippets/stdlib_os.py | 23 +++++++++++++++--------
 vm/src/stdlib/os.rs         |  9 +++++++++
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/tests/snippets/stdlib_os.py b/tests/snippets/stdlib_os.py
index 237acf94f..4bc50ee41 100644
--- a/tests/snippets/stdlib_os.py
+++ b/tests/snippets/stdlib_os.py
@@ -57,6 +57,7 @@ class TestWithTempDir():
 
 FILE_NAME = "test1"
 FILE_NAME2 = "test2"
+FOLDER = "dir1"
 CONTENT = b"testing"
 CONTENT2 = b"rustpython"
 CONTENT3 = b"BOYA"
@@ -75,15 +76,21 @@ with TestWithTempDir() as tmpdir:
 	assert os.read(fd, len(CONTENT3)) == CONTENT3
 	os.close(fd)
 
-
 	fname2 = tmpdir + os.sep + FILE_NAME2
 	with open(fname2, "wb"):
 		pass
-	files = set()
-	paths = set()
-	for dir_entry in os.scandir(tmpdir):
-		files.add(dir_entry.name)
-		paths.add(dir_entry.path)
+	folder = tmpdir + os.sep + FOLDER
+	os.mkdir(folder)
 
-	assert files == set([FILE_NAME, FILE_NAME2])
-	assert paths == set([fname, fname2])
+	names = set()
+	paths = set()
+	dirs = set()
+	for dir_entry in os.scandir(tmpdir):
+		names.add(dir_entry.name)
+		paths.add(dir_entry.path)
+		if dir_entry.is_dir():
+			dirs.add(dir_entry.name)
+
+	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
index 395e269d5..43c65129a 100644
--- 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", {