1

#!/usr/bin/env python

     2

# -*- coding: utf-8 -*-

     3

    24

    25

import os

    26

    27

import wx

    28

    29

from util.BitmapLibrary import GetBitmap

    30

    31

DRIVE, FOLDER, FILE = range(3)

    32

    33

def sort_folder(x, y):

    34

    if x[1] == y[1]:

    35

        return cmp(x[0], y[0])

    36

    elif x[1] != FILE:

    37

        return -1

    38

    else:

    39

        return 1

    40

    41

def splitpath(path):

    42

    head, tail = os.path.split(path)

    43

    if head == "":

    44

        return [tail]

    45

    elif tail == "":

    46

        return splitpath(head)

    47

    return splitpath(head) + [tail]

    48

    49

class FolderTree(wx.Panel):

    50

    51

    def __init__(self, parent, folder, filter=None, editable=True):

    52

        wx.Panel.__init__(self, parent, style=wx.TAB_TRAVERSAL)

    53

    54

        main_sizer = wx.BoxSizer(wx.VERTICAL)

    55

    56

        self.Tree = wx.TreeCtrl(self, 

    57

              style=wx.TR_HAS_BUTTONS|

    58

                    wx.TR_SINGLE|

    59

                    wx.SUNKEN_BORDER|

    60

                    wx.TR_HIDE_ROOT|

    61

                    wx.TR_LINES_AT_ROOT|

    62

                    wx.TR_EDIT_LABELS)

    63

        if wx.Platform == '__WXMSW__':

    64

            self.Bind(wx.EVT_TREE_ITEM_ACTIVATED, self.OnTreeItemExpanded, self.Tree)

    65

            self.Tree.Bind(wx.EVT_LEFT_DOWN, self.OnTreeLeftDown)

    66

        else:

    67

            self.Bind(wx.EVT_TREE_ITEM_EXPANDED, self.OnTreeItemExpanded, self.Tree)

    68

        self.Bind(wx.EVT_TREE_ITEM_COLLAPSED, self.OnTreeItemCollapsed, self.Tree)

    69

        self.Bind(wx.EVT_TREE_BEGIN_LABEL_EDIT, self.OnTreeBeginLabelEdit, self.Tree)

    70

        self.Bind(wx.EVT_TREE_END_LABEL_EDIT, self.OnTreeEndLabelEdit, self.Tree)

    71

        main_sizer.AddWindow(self.Tree, 1, flag=wx.GROW)

    72

    73

        if filter is not None:

    74

            self.Filter = wx.ComboBox(self, style=wx.CB_READONLY)

    75

            self.Bind(wx.EVT_COMBOBOX, self.OnFilterChanged, self.Filter)

    76

            main_sizer.AddWindow(self.Filter, flag=wx.GROW)

    77

        else:

    78

            self.Filter = None

    79

    80

        self.SetSizer(main_sizer)

    81

    82

        self.Folder = folder

    83

        self.Editable = editable

    84

    85

        self.TreeImageList = wx.ImageList(16, 16)

    86

        self.TreeImageDict = {}

    87

        for item_type, bitmap in [(DRIVE, "tree_drive"),

    88

                                  (FOLDER, "tree_folder"),

    89

                                  (FILE, "tree_file")]:

    90

            self.TreeImageDict[item_type] = self.TreeImageList.Add(GetBitmap(bitmap))

    91

        self.Tree.SetImageList(self.TreeImageList)

    92

    93

        self.Filters = {}

    94

        if self.Filter is not None:

    95

            filter_parts = filter.split("|")

    96

            for idx in xrange(0, len(filter_parts), 2):

    97

                if filter_parts[idx + 1] == "*.*":

    98

                    self.Filters[filter_parts[idx]] = ""

    99

                else:

   100

                    self.Filters[filter_parts[idx]] = filter_parts[idx + 1].replace("*", "")

   101

                self.Filter.Append(filter_parts[idx])

   102

                if idx == 0:

   103

                    self.Filter.SetStringSelection(filter_parts[idx])

   104

   105

            self.CurrentFilter = self.Filters[self.Filter.GetStringSelection()]

   106

        else:

   107

            self.CurrentFilter = ""

   108

   109

    def _GetFolderChildren(self, folderpath, recursive=True):

   110

        items = []

   111

        if wx.Platform == '__WXMSW__' and folderpath == "/":

   112

            for c in xrange(ord('a'), ord('z')):

   113

                drive = os.path.join("%s:\\" % chr(c))

   114

                if os.path.exists(drive):

   115

                    items.append((drive, DRIVE, self._GetFolderChildren(drive, False)))

   116

        else:

   117

            try:

   118

                files = os.listdir(folderpath)

   119

            except:

   120

                return []

   121

            for filename in files:

   122

                if not filename.startswith("."):

   123

                    filepath = os.path.join(folderpath, filename)

   124

                    if os.path.isdir(filepath):

   125

                        if recursive:

   126

                            children = len(self._GetFolderChildren(filepath, False))

   127

                        else:

   128

                            children = 0

   129

                        items.append((filename, FOLDER, children))

   130

                    elif (self.CurrentFilter == "" or 

   131

                          os.path.splitext(filename)[1] == self.CurrentFilter):

   132

                        items.append((filename, FILE, None))

   133

        if recursive:

   134

            items.sort(sort_folder)

   135

        return items

   136

   137

    def SetFilter(self, filter):

   138

        self.CurrentFilter = filter

   139

   140

    def GetTreeCtrl(self):

   141

        return self.Tree

   142

   143

    def RefreshTree(self):

   144

        root = self.Tree.GetRootItem()

   145

        if not root.IsOk():

   146

            root = self.Tree.AddRoot("")

   147

        self.GenerateTreeBranch(root, self.Folder)

   148

   149

    def GenerateTreeBranch(self, root, folderpath):

   150

        item, item_cookie = self.Tree.GetFirstChild(root)

   151

        for idx, (filename, item_type, children) in enumerate(self._GetFolderChildren(folderpath)):

   152

            if not item.IsOk():

   153

                item = self.Tree.AppendItem(root, filename, self.TreeImageDict[item_type])

   154

                if wx.Platform != '__WXMSW__':

   155

                    item, item_cookie = self.Tree.GetNextChild(root, item_cookie)

   156

            elif self.Tree.GetItemText(item) != filename:

   157

                item = self.Tree.InsertItemBefore(root, idx, filename, self.TreeImageDict[item_type])

   158

            filepath = os.path.join(folderpath, filename)

   159

            if item_type != FILE:

   160

                if self.Tree.IsExpanded(item):

   161

                    self.GenerateTreeBranch(item, filepath)

   162

                elif children > 0:

   163

                    self.Tree.SetItemHasChildren(item)

   164

            item, item_cookie = self.Tree.GetNextChild(root, item_cookie)

   165

        to_delete = []

   166

        while item.IsOk():

   167

            to_delete.append(item)

   168

            item, item_cookie = self.Tree.GetNextChild(root, item_cookie)

   169

        for item in to_delete:

   170

            self.Tree.Delete(item)

   171

   172

    def ExpandItem(self, item):

   173

        self.GenerateTreeBranch(item, self.GetPath(item))

   174

        self.Tree.Expand(item)

   175

   176

    def OnTreeItemActivated(self, event):

   177

        self.ExpandItem(event.GetItem())

   178

        event.Skip()

   179

   180

    def OnTreeLeftDown(self, event):

   181

        item, flags = self.Tree.HitTest(event.GetPosition())

   182

        if flags & wx.TREE_HITTEST_ONITEMBUTTON and not self.Tree.IsExpanded(item):

   183

            self.ExpandItem(item)

   184

        else:

   185

            event.Skip()

   186

   187

    def OnTreeItemExpanded(self, event):

   188

        item = event.GetItem()

   189

        self.GenerateTreeBranch(item, self.GetPath(item))

   190

        event.Skip()

   191

   192

    def OnTreeItemCollapsed(self, event):

   193

        item = event.GetItem()

   194

        self.Tree.DeleteChildren(item)

   195

        self.Tree.SetItemHasChildren(item)

   196

        event.Skip()

   197

   198

    def OnTreeBeginLabelEdit(self, event):

   199

        item = event.GetItem()

   200

        if self.Editable and not self.Tree.ItemHasChildren(item):

   201

            event.Skip()

   202

        else:

   203

            event.Veto()

   204

   205

    def OnTreeEndLabelEdit(self, event):

   206

        new_name = event.GetLabel()

   207

        if new_name != "":

   208

            old_filepath = self.GetPath(event.GetItem())

   209

            new_filepath = os.path.join(os.path.split(old_filepath)[0], new_name)

   210

            if new_filepath != old_filepath:

   211

                if not os.path.exists(new_filepath):

   212

                    os.rename(old_filepath, new_filepath)

   213

                    event.Skip()

   214

                else:

   215

                    message =  wx.MessageDialog(self, 

   216

                        _("File '%s' already exists!") % new_name, 

   217

                        _("Error"), wx.OK|wx.ICON_ERROR)

   218

                    message.ShowModal()

   219

                    message.Destroy()

   220

                    event.Veto()

   221

        else:

   222

            event.Skip()

   223

   224

    def OnFilterChanged(self, event):

   225

        self.CurrentFilter = self.Filters[self.Filter.GetStringSelection()]

   226

        self.RefreshTree()

   227

        event.Skip()

   228

   229

    def _SelectItem(self, root, parts):

   230

        if len(parts) == 0:

   231

            self.Tree.SelectItem(root)

   232

        else:

   233

            item, item_cookie = self.Tree.GetFirstChild(root)

   234

            while item.IsOk():

   235

                if self.Tree.GetItemText(item) == parts[0]:

   236

                    if (self.Tree.ItemHasChildren(item) and 

   237

                        not self.Tree.IsExpanded(item)):

   238

                        self.Tree.Expand(item)

   239

                        wx.CallAfter(self._SelectItem, item, parts[1:])

   240

                    else:

   241

                        self._SelectItem(item, parts[1:])

   242

                    return

   243

                item, item_cookie = self.Tree.GetNextChild(root, item_cookie)

   244

   245

    def SetPath(self, path):

   246

        if path.startswith(self.Folder):

   247

            root = self.Tree.GetRootItem()

   248

            if root.IsOk():

   249

                relative_path = path.replace(os.path.join(self.Folder, ""), "")

   250

                self._SelectItem(root, splitpath(relative_path))

   251

   252

    def GetPath(self, item=None):

   253

        if item is None:

   254

            item = self.Tree.GetSelection()

   255

        if item.IsOk():

   256

            filepath = self.Tree.GetItemText(item)

   257

            parent = self.Tree.GetItemParent(item)

   258

            while parent.IsOk() and parent != self.Tree.GetRootItem():

   259

                filepath = os.path.join(self.Tree.GetItemText(parent), filepath)

   260

                parent = self.Tree.GetItemParent(parent)

   261

            return os.path.join(self.Folder, filepath)

   262

        return self.Folder