count-packages.hs

#!/usr/bin/env stack
{- stack script
   --resolver lts-9.1
   --package turtle
   --package foldl
   --package system-filepath
   --package bytestring
   --package containers
-}
-- script used to mirror Hackage: https://github.com/AleXoundOS/haskell-stack-mirror-script/issues
-- Usage: ./count-packages.hs <directory with the packages>

{-# LANGUAGE OverloadedStrings #-}
import Turtle
import Prelude hiding (FilePath)
import qualified Control.Foldl as Fold
import qualified Filesystem.Path.CurrentOS as FS
import qualified Data.ByteString.Lazy as BS
import qualified Data.ByteString.Char8 as C8
import qualified Data.Map.Strict as Map

type PkgName = C8.ByteString
type Occurrences = Int
type PkgFile = C8.ByteString

type PkgMap = Map.Map PkgName (PkgFile, Occurrences)

pkgMapInsert :: PkgMap -> (C8.ByteString, C8.ByteString) -> PkgMap
pkgMapInsert m (name, version)
  = case Map.lookup name m of
      Just (pkg, q) -> Map.insert name (max pkg pkg', q + 1) m
      Nothing -> Map.insert name (pkg', 1) m
  where pkg' = name <> version

toPkgMapWith :: (FilePath -> (C8.ByteString, C8.ByteString)) -> [FilePath] -> PkgMap
toPkgMapWith = func Map.empty
  where func m _ [] = m
        func m f (x:xs) = func (pkgMapInsert m x') f xs
          where x' = f x

splitVersion :: FilePath -> (C8.ByteString, C8.ByteString)
splitVersion = C8.breakEnd (== '-') . FS.encode

parseDirectory :: Parser FilePath
parseDirectory = argPath "DIR" "The packages directory"

countFiles :: MonadIO io => FilePath -> io Int
countFiles dir = fold (ls dir) Fold.length

listOfFiles :: MonadIO io => FilePath -> io [FilePath]
listOfFiles dir = fold (ls dir) Fold.list

getLatestVersionsSize :: MonadIO io => PkgMap -> io Size
getLatestVersionsSize m = sum <$> traverse du (getFilePaths m)
    where getFilePaths = map (FS.decode . fst) . Map.elems

sumPackages :: PkgMap -> Int
sumPackages m = sum (getOcurrences m)
  where getOcurrences = map snd . Map.elems

main :: IO ()
main = do
  directory <- options "" parseDirectory
  files <- listOfFiles directory
  -- number of different packages
  let pkgMap = toPkgMapWith splitVersion files
  -- total number of packages
  total <- countFiles directory
  -- total size on disk of the latest version of every package
  latestVersionsSize <- getLatestVersionsSize pkgMap
  totalSize <- sum <$> traverse du files

  printf ("# Packages in this directory: " %w% "\n") total
  printf ("# Different packages in this directory: " %w% "\n") (length pkgMap)
  printf ("# Total size: " %s% "\n") (format sz totalSize)
  printf ("  (only the latest version of each package): " %s% "\n") (format sz latestVersionsSize)