1 files changed, 16 insertions, 8 deletions
diff --git a/app/models/tree.rb b/app/models/tree.rb
index 042050527c1..5fbad19b468 100644
--- a/app/models/tree.rb
+++ b/app/models/tree.rb
@@ -1,17 +1,25 @@
 class Tree
-  attr_accessor :raw
+  attr_accessor :entries, :readme
 
-  def initialize(repository, sha, ref = nil, path = nil)
-    @raw = Gitlab::Git::Tree.new(repository, sha, ref, path)
+  def initialize(repository, sha, path = '/')
+    path = '/' if path.blank?
+    git_repo = repository.raw_repository
+    @entries = Gitlab::Git::Tree.where(git_repo, sha, path)
+
+    if readme_tree = @entries.find(&:readme?)
+      @readme = Gitlab::Git::Blob.find(git_repo, sha, readme_tree.name)
+    end
   end
 
-  def method_missing(m, *args, &block)
-    @raw.send(m, *args, &block)
+  def trees
+    @entries.select(&:dir?)
   end
 
-  def respond_to?(method)
-    return true if @raw.respond_to?(method)
+  def blobs
+    @entries.select(&:file?)
+  end
 
-    super
+  def submodules
+    @entries.select(&:submodule?)
   end
 end