KeywordClassifier.cs

/* Modified version of Thomas Restrepo's KeywordClassifier.cs.
 * The original can be found here:
 * http://github.com/tomasr/KeywordClassifier/blob/b2472cc2bbdd823a4d8a2ff5fdbfc82c4da1423e/KeywordClassifier.cs
 * This has been updated for Beta 2, and switched from using a classifier to a view-specific
 * tagger, to avoid issues around self-consumption.
 */

using System;
using System.Collections.Generic;
using System.ComponentModel.Composition;
using System.Diagnostics;
using System.Linq;
using Microsoft.VisualStudio.Text;
using Microsoft.VisualStudio.Text.Classification;
using Microsoft.VisualStudio.Text.Editor;
using Microsoft.VisualStudio.Text.Tagging;
using Microsoft.VisualStudio.Utilities;

namespace Winterdom.VisualStudio.Extensions.Text
{
    static class Constants
    {
        public const String CLASSIF_NAME = "FlowControl";
        public const String LINQ_CLASSIF_NAME = "LinqOperator";
        public const String VISIBILITY_CLASSIF_NAME = "VisibilityKeyword";
    }

    [Export(typeof(IViewTaggerProvider))]
    [ContentType(CSharp.ContentType)]
    [ContentType(Cpp.ContentType)]
    [TagType(typeof(ClassificationTag))]
    public class KeywordClassifierProvider : IViewTaggerProvider
    {
        [Import]
        internal IClassificationTypeRegistryService ClassificationRegistry = null;

        [Import]
        internal IBufferTagAggregatorFactoryService Aggregator = null;

        public ITagger<T> CreateTagger<T>(ITextView textView, ITextBuffer buffer) where T : ITag
        {
            return new KeywordClassifier(ClassificationRegistry, Aggregator.CreateTagAggregator<ClassificationTag>(buffer)) as ITagger<T>;
        }
    }

    class KeywordClassifier : ITagger<ClassificationTag>
    {
        private IClassificationType keywordClassification;
        private IClassificationType linqClassification;
        private IClassificationType visClassification;
        private ITagAggregator<ClassificationTag> aggregator;

#pragma warning disable 67
        public event EventHandler<SnapshotSpanEventArgs> TagsChanged;
#pragma warning restore 67

        internal KeywordClassifier(
              IClassificationTypeRegistryService registry,
              ITagAggregator<ClassificationTag> aggregator)
        {
            keywordClassification = registry.GetClassificationType(Constants.CLASSIF_NAME);
            linqClassification = registry.GetClassificationType(Constants.LINQ_CLASSIF_NAME);
            visClassification = registry.GetClassificationType(Constants.VISIBILITY_CLASSIF_NAME);
            this.aggregator = aggregator;
        }

        private ILanguageKeywords GetKeywordsByContentType(IContentType contentType)
        {
            if (contentType.IsOfType(CSharp.ContentType))
            {
                return new CSharp();
            }
            else if (contentType.IsOfType(Cpp.ContentType))
            {
                return new Cpp();
            }
            else
            {
                Debug.Fail("How were we called with an invalid content type?");
                return null;
            }
        }

        public IEnumerable<ITagSpan<ClassificationTag>> GetTags(NormalizedSnapshotSpanCollection spans)
        {
            if (spans.Count == 0)
            {
                yield break;
            }

            ITextSnapshot snapshot = spans[0].Snapshot;

            ILanguageKeywords keywords = GetKeywordsByContentType(snapshot.TextBuffer.ContentType);
            if (keywords == null)
            {
                yield break;
            }

            // find spans that the language service has already classified as keywords ...
            foreach (var classifiedSpan in from cs in aggregator.GetTags(spans)
                                           let name = cs.Tag.ClassificationType.Classification.ToLower()
                                           where name.Contains("keyword")
                                           let classifiedSpans = cs.Span.GetSpans(snapshot)
                                           where classifiedSpans.Count > 0
                                           select classifiedSpans[0])
            {
                string text = classifiedSpan.GetText();

                if (keywords.ControlFlow.Contains(text))
                    yield return new TagSpan<ClassificationTag>(classifiedSpan, new ClassificationTag(keywordClassification));
                else if (keywords.Linq.Contains(text))
                    yield return new TagSpan<ClassificationTag>(classifiedSpan, new ClassificationTag(linqClassification));
                else if (keywords.Visibility.Contains(text))
                    yield return new TagSpan<ClassificationTag>(classifiedSpan, new ClassificationTag(visClassification));
            }
        }
    }
}