treemap-html-tools-0.1: src/dirmap.hs

import Data.List
import Data.Maybe
import Control.Monad 
import System.FilePath 
import System.Environment
import System.Directory
import Data.Tree
import Text.HTML.TreeMap
import System.Console.GetOpt

maxNameLen = 15

data LsDirOpts = LsDirOpts { lsMaxName :: Int
                           , lsMaxDep  :: Maybe Int
                           , lsFiles   :: Bool
                           , lsEmptyDir :: Bool }

defLsOpts :: LsDirOpts
defLsOpts = LsDirOpts maxNameLen Nothing False True

data Flag = Verbose  
          | Level String
          | MaxName String
          | Files
          | NoEmptyDir
          | Output String deriving Show

options :: [OptDescr Flag]
options =
    [ Option ['v']     ["verbose"] (NoArg Verbose)                       "verbose output on stderr"
    , Option ['l']     ["level"]   (ReqArg (Level)  "Integer")           "maxlevel for recursive directory listing"
    , Option ['n']     ["name"]    (ReqArg (MaxName) "Integer")          "maximum size of filename. than truncation."
    , Option ['o']     ["output"]  (ReqArg (Output) "FILE")              "output FILE"
    , Option ['f']     ["files"]   (NoArg Files)                         "include files in listing"
    , Option ['e']     []          (NoArg NoEmptyDir)                    "do not included subtrees without files"
    ]

checkOpts :: [String] -> IO ([Flag], [String])
checkOpts argv = 
    case getOpt Permute options argv of
      (o,[],[] ) -> ioError (userError (usageInfo header options))
      (o,n,[]  ) -> return (o,n)
      (_,_,errs) -> ioError (userError (concat errs ++ usageInfo header options))
    where header = "Usage: dirmap [OPTION...] directory"

lsOpts :: [Flag] -> LsDirOpts
lsOpts xs = foldr (procFlag) defLsOpts xs 
    where
      procFlag (Level l)  o = if null l then o else o { lsMaxDep = Just (read l :: Int) }
      procFlag Files      o = o { lsFiles = True }
      procFlag (MaxName n) o = o { lsMaxName = (read n :: Int) }
      procFlag NoEmptyDir o = o { lsEmptyDir = False }
      procFlag _          o = o

filename :: [Flag] -> FilePath
filename []             = "dirmap.html"
filename ((Output d):_) = d
filename (_:xs)         = filename xs

main = do
  args <- getArgs
  (flags,names) <- checkOpts args
  print flags
  print names
  dirTree <- lsDirTree (lsOpts flags) (head names)
  putStrLn $ drawTree $ dirTree
  writeFile (filename flags) (treeMap dirTree)
  putStrLn $ "Written: "++(filename flags)

lsDirTree :: LsDirOpts -> String -> IO (Tree String)
lsDirTree opts@(LsDirOpts _ (Just 0) _ _) dir = return $ Node ((shortName opts dir) </> "...") []
lsDirTree opts dir = do 
  entries <- liftM (filterDots) (getDirectoryContents dir)
  dirEnts <- filterM (\x -> doesDirectoryExist (dir </> x)) entries 
  fileEnts <- filterM (\x -> liftM not $ doesDirectoryExist (dir </> x)) entries
  let files = if (lsFiles opts) then fileEnts else []
  subTrees <- mapM (\x -> lsDirTree (decDep opts) (dir </> x)) dirEnts
  let subTreesM = if (lsEmptyDir opts) then subTrees else filter hasFiles subTrees 
  let dirNode = Node ((shortName opts dir) ++ "/") (subTreesM ++ (map (\x -> Node (shortName opts x) []) files))
  return dirNode

hasFiles :: Tree String -> Bool
hasFiles xs = let ys = flatten xs
                  isFile [] = False
                  isFile x = (last x) /= '/'
              in any (isFile) (concatMap lines ys)

shortName opts x = let name = last $ splitPath x
                       maxNL = lsMaxName opts
                   in if length name > maxNL then (take maxNL name) ++ "..." else name

filterDots = filter (\e -> e /="." && e /= "..")

decDep :: LsDirOpts -> LsDirOpts
decDep opts@(LsDirOpts _ Nothing _ _)      = opts
decDep opts@(LsDirOpts _ (Just (x+1)) _ _) = opts { lsMaxDep = Just x }