Add --prune to stack dot

ChangeLog.md

@@ -1,6 +1,8 @@
 ## Unreleased
 
-* Add --ignore-subdirs flag to init command [#435](https://github.com/commercialhaskell/stack/pull/435)
+* Add `--prune` flag to `stack dot` [#487](https://github.com/commercialhaskell/stack/issues/487)
+* Add `--[no-]external`,`--[no-]include-base` flags to `stack dot` [#437](https://github.com/commercialhaskell/stack/issues/437)
+* Add `--ignore-subdirs` flag to init command [#435](https://github.com/commercialhaskell/stack/pull/435)
 * Handle attempt to use non-existing resolver [#436](https://github.com/commercialhaskell/stack/pull/436)
 * Add `--force` flag to `init` command
 * exec style commands accept the `--package` option (see [Reddit discussion](http://www.reddit.com/r/haskell/comments/3bd66h/stack_runghc_turtle_as_haskell_script_solution/))

src/Stack/Dot.hs

@@ -7,6 +7,7 @@ module Stack.Dot (dot
                  ,dotOptsParser
                  ,resolveDependencies
                  ,printGraph
+                 ,pruneGraph
                  ) where
 
 import           Control.Monad (void)
@@ -15,8 +16,10 @@ import           Control.Monad.IO.Class
 import           Control.Monad.Logger (MonadLogger, logInfo)
 import           Control.Monad.Reader (MonadReader)
 import           Control.Monad.Trans.Control (MonadBaseControl)
+import           Data.Char (isSpace)
 import qualified Data.Foldable as F
 import qualified Data.HashSet as HashSet
+import           Data.List.Split (splitOn)
 import           Data.Map (Map)
 import qualified Data.Map as Map
 import           Data.Set (Set)
@@ -42,11 +45,17 @@ data DotOpts = DotOpts
     -- ^ Include dependencies on base
     , dotDependencyDepth :: Maybe Int
     -- ^ Limit the depth of dependency resolution to (Just n) or continue until fixpoint
+    , dotPrune :: Set String
+    -- ^ Package names to prune from the graph
     }
 
 -- | Parser for arguments to `stack dot`
 dotOptsParser :: Parser DotOpts
-dotOptsParser = DotOpts <$> includeExternal <*> includeBase <*> depthLimit
+dotOptsParser = DotOpts
+            <$> includeExternal
+            <*> includeBase
+            <*> depthLimit
+            <*> fmap (maybe Set.empty Set.fromList . fmap splitNames) prunedPkgs
   where includeExternal = boolFlags False
                                     "external"
                                     "inclusion of external dependencies"
@@ -61,6 +70,15 @@ dotOptsParser = DotOpts <$> includeExternal <*> includeBase <*> depthLimit
                               metavar "DEPTH" <>
                               help ("Limit the depth of dependency resolution " <>
                                     "(Default: No limit)")))
+        prunedPkgs = optional (strOption
+                                   (long "prune" <>
+                                    metavar "PACKAGES" <>
+                                    help ("Prune each package name " <>
+                                          "from the comma separated list " <>
+                                          "of package names PACKAGES")))
+
+        splitNames :: String -> [String]
+        splitNames = map (takeWhile (not . isSpace) . dropWhile isSpace) . splitOn ","
 
 -- | Visualize the project's dependencies as a graphviz graph
 dot :: (HasEnvConfig env
@@ -81,15 +99,46 @@ dot dotOpts = do
     resultGraph <- withLoadPackage menv (\loader -> do
       let depLoader = createDepLoader sourceMap (fmap3 packageAllDeps loader)
       liftIO $ resolveDependencies (dotDependencyDepth dotOpts) graph depLoader)
-    printGraph dotOpts locals (if dotIncludeBase dotOpts
-                                 then resultGraph
-                                 else filterOutDepsOnBase resultGraph)
-  where filterOutDepsOnBase = Map.filterWithKey (\k _ -> show k /= "base") .
-                              fmap (Set.filter ((/= "base") . show))
-        -- fmap a function over the result of a function with 3 arguments
-        fmap3 :: Functor f => (d -> e) -> (a -> b -> c -> f d) -> (a -> b -> c -> f e)
+    let pkgsToPrune = if dotIncludeBase dotOpts
+                         then dotPrune dotOpts
+                         else Set.insert "base" (dotPrune dotOpts)
+        localNames = Set.fromList (map (packageName . lpPackage) locals)
+        prunedGraph = pruneGraph localNames pkgsToPrune resultGraph
+    printGraph dotOpts locals prunedGraph
+  where -- fmap a function over the result of a function with 3 arguments
+        fmap3 :: Functor f => (d -> e) -> (a -> b -> c -> f d) -> a -> b -> c -> f e
         fmap3 f g a b c = f <$> g a b c
 
+-- | `pruneGraph dontPrune toPrune graph` prunes all packages in
+-- `graph` with a name in `toPrune` and removes resulting orphans
+-- unless they are in `dontPrune`
+pruneGraph :: (F.Foldable f, F.Foldable g)
+           => f PackageName
+           -> g String
+           -> Map PackageName (Set PackageName)
+           -> Map PackageName (Set PackageName)
+pruneGraph dontPrune names =
+  pruneUnreachable dontPrune . Map.mapMaybeWithKey (\pkg pkgDeps ->
+    if show pkg `F.elem` names
+      then Nothing
+      else let filtered = Set.filter (\n -> show n `F.notElem` names) pkgDeps
+           in if Set.null filtered && not (Set.null pkgDeps)
+                then Nothing
+                else Just filtered)
+
+-- | Make sure that all unreachable nodes (orphans) are pruned
+pruneUnreachable :: F.Foldable f
+                 => f PackageName
+                 -> Map PackageName (Set PackageName)
+                 -> Map PackageName (Set PackageName)
+pruneUnreachable dontPrune = fixpoint prune
+  where fixpoint :: Eq a => (a -> a) -> a -> a
+        fixpoint f v = if f v == v then v else fixpoint f (f v)
+        prune graph' = Map.filterWithKey (\k _ -> reachable k) graph'
+          where reachable k = k `F.elem` dontPrune || k `Set.member` reachables
+                reachables = F.fold graph'
+
+
 -- | Resolve the dependency graph up to (Just depth) or until fixpoint is reached
 resolveDependencies :: (Applicative m, Monad m)
                     => Maybe Int
@@ -138,10 +187,12 @@ printGraph :: (Applicative m, MonadLogger m)
            -> m ()
 printGraph dotOpts locals graph = do
   $logInfo "strict digraph deps {"
-  printLocalNodes dotOpts locals
+  printLocalNodes dotOpts filteredLocals
   printLeaves graph
   void (Map.traverseWithKey printEdges graph)
   $logInfo "}"
+  where filteredLocals = filter (\local ->
+          show (packageName (lpPackage local)) `Set.notMember` dotPrune dotOpts) locals
 
 -- | Print the local nodes with a different style depending on options
 printLocalNodes :: (F.Foldable t, MonadLogger m)

src/test/Stack/DotSpec.hs

@@ -3,14 +3,18 @@
 module Stack.DotSpec where
 
 import           Data.ByteString.Char8 (ByteString)
+import qualified Data.Foldable as F
 import           Data.Functor.Identity
+import           Data.List ((\\))
 import qualified Data.Map as Map
 import           Data.Maybe (fromMaybe)
 import           Data.Set (Set)
 import qualified Data.Set as Set
 import           Options.Applicative (execParserPure,idm,prefs,info,getParseResult)
 import           Stack.Types
 import           Test.Hspec
+import           Test.Hspec.QuickCheck (prop)
+import           Test.QuickCheck (forAll,sublistOf)
 
 import           Stack.Dot
 
@@ -38,6 +42,21 @@ spec = do
            resultGraph' = resolveDependencies Nothing graph' stubLoader
        fmap Map.size resultGraph' `shouldBe` fmap ((+1) . Map.size) resultGraph
 
+    prop "requested packages are pruned" $ do
+      let resolvedGraph = runIdentity (resolveDependencies Nothing graph stubLoader)
+          allPackages g = Set.map show (Map.keysSet g `Set.union` F.fold g)
+      forAll (sublistOf (Set.toList (allPackages resolvedGraph))) $ \toPrune ->
+        let pruned = pruneGraph [pkgName "one", pkgName "two"] toPrune resolvedGraph
+        in Set.null (allPackages pruned `Set.intersection` Set.fromList toPrune)
+
+    prop "pruning removes orhpans" $ do
+      let resolvedGraph = runIdentity (resolveDependencies Nothing graph stubLoader)
+          allPackages g = Set.map show (Map.keysSet g `Set.union` F.fold g)
+          orphans g = Map.filterWithKey (\k _ -> not (graphElem k g)) g
+      forAll (sublistOf (Set.toList (allPackages resolvedGraph))) $ \toPrune ->
+        let pruned = pruneGraph [pkgName "one", pkgName "two"] toPrune resolvedGraph
+        in null (Map.keys (orphans pruned) \\ [pkgName "one", pkgName "two"])
+
   where graphElem e graph = Set.member e . Set.unions . Map.elems $ graph
 
 {- Helper functions below -}
@@ -46,7 +65,7 @@ spec = do
 pkgName :: ByteString -> PackageName
 pkgName = fromMaybe failure . parsePackageName
   where
-   failure = (error "Internal error during package name creation in DotSpec.pkgName")
+   failure = error "Internal error during package name creation in DotSpec.pkgName"
 
 -- Stub, simulates the function to load package dependecies
 stubLoader :: PackageName -> Identity (Set PackageName)

stack.cabal

@@ -141,6 +141,7 @@ library
                    , process >= 1.2.0.0
                    , resourcet >= 1.1.4.1
                    , safe >= 0.3
+                   , split
                    , stm >= 2.4.4
                    , streaming-commons >= 0.1.10.0
                    , tar >= 0.4.1.0
@@ -204,6 +205,7 @@ test-suite stack-test
                 , Stack.BuildPlanSpec
                 , Stack.Build.ExecuteSpec
                 , Stack.ConfigSpec
+                , Stack.DotSpec
                 , Stack.PackageDumpSpec
                 , Stack.ArgsSpec
                 , Network.HTTP.Download.VerifiedSpec
@@ -228,6 +230,7 @@ test-suite stack-test
                 , text
                 , optparse-applicative
                 , bytestring
+                , QuickCheck
   default-language:    Haskell2010
 
 test-suite stack-integration-test