Dirsum.hs

module Main where

-- A simple directory summing thing for Haskell.
-- by Daniel Lyons <fusion@storytotell.org>
-- BSD licensed.

import Control.Applicative

import Data.Either
import Data.Foldable
import Data.Function
import Data.List
import Data.Maybe
import Data.Monoid
import Data.Tree

import System.Directory
import System.Environment
import System.FilePath
import System.IO
import System.IO.Error

import Text.Printf

-- alias for byte count
type Size = Integer

-- the type for holding a filename and its size
data File = File FilePath Size

instance Show File where
    show (File path size) = path ++ " (" ++ (showSize size) ++ ")"

getSize (File _ size) = size
getPath (File path _) = path

showSize :: Size -> String
showSize size' = printf "%0.3f %s" fix mag
    where
      size = fromIntegral size'
      fix = size / (1024 ^ n) :: Float
      magnitudes = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]
      n = min (fromEnum $ log size / log 1024) ((length magnitudes) - 1)
      mag = magnitudes !! n

-- build part of the file tree
buildUnderPath :: FilePath -> IO (File, [FilePath])
-- buildUnderPath path = (,) <$> getFile path <*> safeDirectoryContents path
buildUnderPath path = do
  file <- getFile path
  contents <- safeDirectoryContents path
  return $ (file, contents)


getFile :: FilePath -> IO File
getFile path = File path <$> safeGetFileSize path

safeGetFileSize :: FilePath -> IO Integer
safeGetFileSize path = do
  res <- try $ withFile path ReadMode hFileSize
  return $ either (const 0) id res

safeDirectoryContents :: FilePath -> IO [FilePath]
safeDirectoryContents path = do
 -- get Either the directory contents or an exception
  res <- try $ getDirectoryContents path
  -- return empty if there was an exception or else the contents minus . and ..
  return $ either (const []) (map (path </>) . drop 2) res

-- build the entire file tree under some path
examinePath :: FilePath -> IO (Tree File)
examinePath = unfoldTreeM buildUnderPath

-- calculate the total size of the tree
calculateSize :: Tree File -> Size
calculateSize = getSum . fold . fmap (Sum . getSize)

drawTreeS :: (Show a) => Tree a -> String
drawTreeS = drawTree . fmap show

putTree :: (Show a) => Tree a -> IO ()
putTree = putStr . drawTreeS

main = do
  args <- getArgs
  forM_ args (\path -> do
         tree <- examinePath path
         putTree tree
         putStrLn $ "Total size: " ++ showSize (calculateSize tree))